F ED BN: F EDERATED L EARNING ON N ON -IID F EATURES VIA L OCAL B ATCH N ORMALIZATION
A BSTRACT
The emerging paradigm of federated learning (FL) strives to enable collaborative training of deep models on the network edge without centrally aggregating raw data and hence improving data privacy.
In most cases, the assumption of independent and identically distributed samples across local clients does not hold for federated learning setups.
Under this setting, neural network training performance may vary signiﬁcantly according to the data distribution and even hurt training convergence.
Most of the previous work has focused on a difference in the distribution of labels or client shift.
Unlike those settings, we address an important problem of FL, e.g., different scanner/sensors in medical imaging, different scenery distribution in autonomous driving (highway vs. city), where local clients store examples with different distributions compared to other clients, which we denote as feature shift non-iid .
In this work, we propose an effective method that uses local batch normalization to alleviate the feature shift before averaging models.
The resulting scheme, called FedBN, outperforms both classical FedAvg, as well as the state-of-the-art for non-iid data (FedProx) on our extensive experiments.
These empirical results are supported by a convergence analysis that shows in a simpliﬁed setting that FedBN has a faster convergence rate than FedAvg. Code is available at https://github.com/med-air/FedBN.
I NTRODUCTION
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 signiﬁcantly 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
[Figure] [Figure]
in the feature space, which has not yet been explored in the literature.
Speciﬁcally 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 for all 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 2 with y = cos( w true x ) + (cid:15) , where x ∈ R is drawn iid from a Gaussian distribution and (cid:15) is zero-mean Gaussian noise and consider models of the form f w ( 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 simpliﬁed 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 beneﬁcial.
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.
The ﬁgure 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 beneﬁt shown in the toy example, we also show the beneﬁts 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 signiﬁcantly practical improvements on the extensive experiments.
R ELATED W ORK
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 “ﬂat” 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, there are so far 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 afﬁne distribution shift and tackles this problem by learning the afﬁne transformation.
This hampers the generalization when we cannot estimate the explicit afﬁne 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 beneﬁts given by batch normalization.
Santurkar et al.
(2018) showed that BN makes the optimization landscape signiﬁcantly 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.
P RELIMINARY
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.
Hsieh et al.
(2019) can be described by the joint probability between features x and labels y on each client.
We can rewrite P i ( x , y ) as P i
( y | x ) P i ( x ) and P i ( x | y ) P
i ( y ) .
We deﬁne feature shift as the case that covers: 1) covariate shift : the marginal distributions P
i
( x ) varies across clients, even if P i ( y | x ) is the same for all client; and 2) concept shift : the conditional distribution P i ( 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).
F EDERATED A VERAGING WITH L OCAL B ATCH N ORMALIZATION 4. P ROPOSED M ETHOD - F ED BN
We propose an efﬁcient 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 modiﬁcation results in signiﬁcant empirical improvements in non-iid settings.
We provide an explanation for these improvements in a simpliﬁed scenario, in which we show that FedBN improves the convergence rate under feature shift.
4. P ROBLEM S ETUP
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 { ( x i j , y i j
) ∈ R d × R : j ∈ [ M ] } .
Furthermore, we assume a two-layer neural network with ReLU activations trained by gradient descent.
Let v k ∈ R d denote the parameters of the ﬁrst layer, where k ∈ [ m ] and m is the width of the hidden layer.
Let (cid:107) v (cid:107) S (cid:44) √ v (cid:62)
Sv denote the induced vector norm for a positive deﬁnite 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.
; Li et al.
To be more precise, we make the following assumption.
Assumption 4.1 (Data Distribution) .
For each client i ∈ [ N ] the inputs x ij are centered ( E x i = 0 ) with covariance matrix S i
=
E x i x i (cid:62) , where S i is independent from the label y and may differ for each i ∈ [ N ] e.g., S i are not all identity matrices, and for each index pair p (cid:54)
=
q , x p (cid:54) = κ · x q for all κ ∈ R \ { 0 } .
With Assumption 4.1, the normalization of the ﬁrst layer for client i is v (cid:62) k x i (cid:107) v k (cid:107) S i .
FedBN with clientspeciﬁed BN parameters trains a model f ∗ : R d → R parameterized by ( V , γ, c ) ∈ R m × d ×
[Equation]
m
[Equation]
, i.e.,
[Equation]
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 : R d → 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:
[Equation]
where α 2 controls the magnitude of v k 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 :
[Equation]
4. C ONVERGENCE A NALYSIS
Here we study the trajectory of networks FedAvg ( f ) and FedBN ( f ∗ )’s prediction through the neural tangent kernel (NTK) introduced by Jacot et al.
Recent machine learning theory studies (Arora et al., 2019; Du et al., 2018; Allen-Zhu et al., 2019; van den Brand et al., 2020;
Dukler et al., 2020) have shown that for ﬁnite-width over-parameterized networks, the convergence rate is controlled by the least eigenvalue of the induced kernel in the training evolution.
To simplify tracing the optimization dynamics, we consider the case that the number of local updates
E is 1 .
We can decompose the NTK into a magnitude component G ( t ) and direction component
V ( t ) /α
2 following Dukler et al.
[Equation]
where Λ (
[Equation]
The speciﬁc forms of V ( t ) and G ( t ) are given in Appendix B.1.
Let λ min ( A ) denote the minimal eigenvalue of matrix A .
The matrices V ( t ) and G ( t ) are positive semi-deﬁnite, since they can be viewed as covariance matrices.
This gives λ min ( Λ ( t )) ≥ max (cid:8) λ min ( V ( t )) /α 2 , λ min ( G ( t )) (cid:9) .
According to NTK, the convergence rate is controlled by λ min ( Λ ( t )) .
Then, for α > 1 , convergence is dominated by G ( t ) .
Let Λ ( t ) and Λ ∗ ( t ) denote the evolution dynamics of FedAvg and FedBN and let G ( t ) and G ∗ ( t ) denote the magnitude component in the evolution dynamics of FedAvg and
FedBN.
For the convergence analysis, we use the auxiliary version of the Gram matrices, which is deﬁned as follows.
Deﬁnition 4.2.
Given sample points { x p } NM p =1 , we deﬁne the auxiliary Gram matrices G ∞ ∈ R NM × NM and G ∗∞ ∈ R NM × NM as
[Equation]
Given Assumption 4.1, we use the key results in Dukler et al.
(2020) to show that G ∞ is positive deﬁnite.
Further, we show that G ∗∞ is positive deﬁnite.
We use the fact that the distance between
G ( t ) and its auxiliary version is small in over-parameterized neural network, such that G ( t ) remains positive deﬁnite.
Lemma 4.3.
Fix points { x p } NMp =
satisfying Assumption 4.1.
Then Gram matrices G ∞ and G ∗∞ deﬁned as in (4) and (5) are strictly positive deﬁnite.
Let the least eigenvalues be λ min ( G ∞ )
=:
µand λ min ( G ∗∞ ) =: µ ∗ 0 , where µ 0 , µ ∗ 0 > 0 .
Proof sketch The main idea follows Du et al.
Dukler et al.
that given points { x p } NMp =1 , the matrices G ∞ and G ∗∞ can be shown as covariance matrix of linearly independent operators.
More details of the proof are given in the Appendix B.2.
Based on our formulation, the convergence rate of FedAvg (Theorem 4.4) can be derived from
Dukler et al.
by considering non-identical covariance matrices.
We derive the convergence rate of FedBN in Corollary 4.5.
Our key result of comparing the convergence rates between FedAvg and FedBN is culminated in Corollary 4.6.
Theorem 4.4 ( G -dominated convergence for FedAvg Dukler et al.
Suppose network (4) is initialized as in (2) with α > 1 , trained using gradient descent and Assumptions 4.1 holds.
Given the loss function of training the neural network is the square loss with targets y satisfying (cid:107) y (cid:107) ∞
=
O (1) .
If m = Ω (cid:0) max (cid:8) N 4 M 4 log( NM/δ ) /α 4 µ 40 , N 2 M 2 log( NM/δ ) /µ 20 (cid:9)(cid:1) , then with probability 1 − δ ,
For iterations t = 0 , 1 , · · · , the evolution matrix Λ ( t ) satisﬁes λ [Equation]
Training with gradient descent of step-size η = O (cid:16)
[Equation]
converges linearly as
[Equation]
Following the key ideas in Dukler et al.
, here we further characterize the convergence for
FedBN.
Corollary 4.5 ( G -dominated convergence for FedBN) .
Suppose network (5) and all other conditions in Theorem 4.4.
With probability 1 − δ , for iterations t = 0 , 1 , · · · , the evolution matrix Λ ∗ ( t ) satisﬁes λ min ( Λ ∗ ( t )) ≥ µ ∗ 0 2 and training with gradient descent of step-size η = O (cid:16) [Equation] converges linearly as (cid:107) f ∗ ( t ) − y (cid:107) 22 ≤ (cid:16) 1 − ηµ
[Equation] 2 (cid:17) t (cid:107) f ∗ (0) − y (cid:107) [Equation]
[Figure]
The exponential factor of convergence for FedAvg (1 − ηµ 0 / 2) and FedBN (1 − ηµ ∗ 0 / 2) are controlled by the smallest eigenvalue of G ( t ) , respectively G ∗ ( t ) .
Then we can analyze the convergence performance of FedAvg and FedBN by comparing λ min ( G ∞ ) and λ min ( G ∗∞ ) .
Corollary 4.6 (Convergence rate comparison between FedAvg and FedBN) .
For the G -dominated convergence, the convergence rate of FedBN is faster than that of FedAvg.
Proof sketch The key is to show λ min ( G ∞ ) ≤ λ min ( G ∗∞ ) .
Comparing equation (4) and (5),
G ∗∞ takes the M × M block matrices on the diagonal of G ∞ .
Let G ∞ i be the i -th M × M block matrices on the diagonal of G ∞ .
By linear algebra, λ min ( G ∞ i ) ≥ λ min ( G ∞ ) for i ∈
[ N ] .
Since G ∗∞ = diag ( G ∞ 1 , · · · , G ∞ N ) , we have λ min ( G ∗∞ ) =
min i ∈
[ N ] { λ min ( G ∞ i ) } .
Therefore, we have the result λ min ( G ∗∞ ) ≥ λ min ( G ∞ ) .
E XPERIMENTS
In this section, we demonstrate that using local BN parameters is beneﬁcial in the presence of feature shift across clients with heterogeneity data.
Our novel local parameter sharing strategy, FedBN, achieves more robust and faster convergence for feature shift non-iid datasets and obtains better model performance compared to alternative methods.
This is shown on both benchmark and large real-world datasets.
5. B ENCHMARK E XPERIMENTS
Settings: We perform an extensive empirical analysis using a benchmark digits classiﬁcation task containing different data sources with feature shift where each dataset is from a different domain.
Data of different domains have heterogeneous appearance but share the same labels and label distribution.
Speciﬁcally, we use the following ﬁve datasets: SVHN Netzer et al.
(2011), USPS Hull (1994), SynthDigits Ganin & Lempitsky (2015), MNIST-M Ganin & Lempitsky (2015) and MNIST LeCun et al.
To match the setup in Section 4, we truncate the sample size of the ﬁve datasets to their smallest number with random sampling, resulting in 7438 training samples in each dataset.
Testing samples are held out and kept the same for all the experiments on this benchmark dataset.
Our classiﬁcation model is a convolutional neural network where BN layers are added following each feature extraction layer (i.e., both convolutional and fully-connected). The architecture is detailed in Appendix D.2.
For model training, we use the cross-entropy loss and SGD optimizer with a learning rate of 10 − 2 .
If not speciﬁed, our default setting for local update epochs is E = 1 , and the default setting for the amount of data at each client is 10% 3 of the dataset original size.
For the default non-iid setting, FL system contains ﬁve clients. Each client exclusively owns data sampled from one of the ﬁve datasets.
More details are listed in Appendix D.2.
[Figure]
Overviews: In the following paragraphs, we present a comprehensive investigation on the properties of the proposed FedBN approach, including: (1) convergence rate; (2) behavior with respect to the choices of local update epochs; (3) performance on various amounts of data at each client; (4) effects at different level of heterogeneity; (5) comparison to state of the art (FedProx (Li et al., 2020b)), and two baselines (FedAvg and SingleSet, i.e., training an individual model within each client). In Appendix G, we also provide empirical results for including a new client with data from an unknown domain into the learning system.
Convergence Rate: We analyze the training loss curve of FedBN in comparison with FedAvg, as shown in Fig.
The loss of FedBN goes down faster and smoother than FedAvg, indicating that
FedBN has a larger convergence rate.
Moreover, compared to FedAvg, FedBN presents smoother and more stable loss curves during learning.
These experimental observations show consensus with what given by Corollary 4.6.
In addition, we present a more comprehensive comparison with different local update epochs E on convergence rate of FedBN and FedAvg (see Appendix E.1).
The results show similar patterns as in Fig 3.
Analysis of Local Updating Epochs: Aggregating at different frequencies may affect the learning behaviour.
Although our theory and the default setting for the other experiment takes E = 1 , we demonstrate FedBN is effective for cases when E > 1 .
In Fig.4 (a), we explore E = 1 , 4 , 8 ,and compare FedBN to baseline FedAvg.
As expected, an inverse relationship between the local updating epochs E and testing accuracy implied for both FedBN and FedAvg shown in Fig.4 (a).
Zooming into the ﬁnal testing accuracy, FedBN reaches stable accuracy ( 0 . 840 ± 0 . 015 ) on various
E and exceeds the accuracy of FedAvg.
Analysis of Local Dataset Size: We vary the data amount for each client from 100% to 1% of its original dataset size, in order to observe FedBN behaviour over different data capacities at each client.
The results in Fig.4 (b) present the accuracy of FedBN and SingleSet 4 .
Testing accuracy starts to signiﬁcantly drop when each of the local client is only attributed 20% percentage of data from its original data amount.
The improvement margin gained from FedBN increases as local dataset sizes decrease.
The results indicate that FedBN can effectively beneﬁt from collaborative training on distributed data, especially when each client only holds a small amount of data which are non-iid.
Effects of Statistical Heterogeneity: A salient question that arises is: to what degree of heterogeneity on feature shift FedBN is superior to FedAvg.
To answer the question, we simulate a federated settings with varying heterogeneity as described below.
We parcel each dataset intosubsets, one for each clients, with equal number of data samples and the same label distribution.
We treat the clients generated from the same dataset as iid clients, while the clients generated from different datasets as non-iid clients.
[Table]
We start with including one client from each dataset in FL system.
Then, we simultaneously add one client from each datasets while keep the existing clients n times, for n ∈ { 1 , . .
For each setting, we train models from scratch.
More clients correspond to less heterogenity.
We show the testing accuracy under different level of heterogeneity in Fig.
4 (c) and include a comparison with FedAvg, which is designed for iid FL.
Our FedBN achieves substantially higher testing than FedAvg over all levels of heterogeneity.
Comparison with state-of-theart: To further validate our method, we compare FedBN with one of the current state-of-the-art methods for non-iid FL, FedProx Li et al.
(2020b), which also shares the beneﬁt of easy adaptation to current FL frameworks in practice.
We also include training on SingleSet and FedAvg as baselines.
For each strategy, we split an independent testing datasets on clients and report the accuracy on the testing datasets.
We perform 5-trial repeating experiment with different random seeds.
The mean and standard deviation of the accuracy on each dataset over trials are shown in Fig.
From the results, we can make the following observation: (1) FedBN achieves the highest accuracy, consistently outperforming the state-of-the-art and baseline methods; (2) FedBN achieves the most signiﬁcant improvements on SVHN whose image appearance is very different from others (i.e., presenting more obvious feature shift); (3) FedBN shows a smaller variance in error over multiple runs, indicating its stability.
5. E XPERIMENTS ON REAL - WORLD DATASETS
To better understand how our proposed algorithm can be beneﬁcial in real-word feature-shift noniid, we have extensively validated the effectiveness of FedBN in comparison with other methods on three real-world datasets: image classiﬁcation on Ofﬁce-Caltech
(Gong et al., 2012) with images acquired in different cameras or environments; image classiﬁcation on DomainNet (Peng et al., 2019) with different image styles; and a neurodisorder diagnosis task on ABIDE I (Di Martino et al., 2014) with patients from different medical institutions 7 .
Datasets and Setup: (1) We conduct the classiﬁcation task on natural images from OfﬁceCaltech10 , which has four data sources composing Ofﬁce-31 Saenko et al.
(2010) (three data sources) and Caltech-256 datasets (one data source)
Grifﬁn et al.
(2007), which are acquired using different camera devices or in different real environment with various background.
Each client joining the FL system is assigned data from one of the four data sources.
Thus data is non-iid across the clients.
(2) Our second dataset is DomainNet , which contains natural images coming from six different data sources: Clipart, Infograph, Painting, Quickdraw, Real, and Sketch.
Similar to (1), each client contains iid data from one of the data sources, but clients with different data sources have different feature distributions.
(3) We include four medical institutions (NYU, USM, UM,
UCLA; each is viewed as a client) from ABIDE I that collects functional brain images using different imaging equipment and protocols.
We validate on a medical application for binary classiﬁcation between autism spectrum disorders patients and healthy control subjects.
The Ofﬁce-Caltech10 contains ten categories of objects.
The DomainNet extensively containsobject categories and we use the top ten most common classes to form a sub-dataset for our experiments.
Our classiﬁcation models adopt AlexNet (Krizhevsky et al., 2012) architecture with BN added after each convolution and fully-connected layer.
Before feeding into the network, all images are resized to 256 × 256 × 3 .
For ABIDE I, each instance is represented as a 5995-dimensional vector through brain connectome computation.
We use a three-layer fully connected neural network as the classiﬁer with the hidden layers of 16 with two BN layers after the ﬁrst two fully connected layers.
Same as the above benchmark, we perform 5 repeated runs for each experiment.
Results and Analysis: The experimental results are shown in Table 1 in the form of mean (std).
On Ofﬁce-Caltech10, FedBN signiﬁcantly outperforms the state-of-the-art method of FedProx, and improves at least 6% on mean accuracy compared with all the alternative methods.
On DomainNet,
FedBN achieved supreme accuracy over most of the datasets.
Interestingly, we ﬁnd the alternative
FL methods achieves comparable results with SingleSet except Quickdraw, and FedBN outperforms them over 10% .
Surprisingly, for the above two tasks, the alternative FL strategies are ineffective in the feature shift non-iid datasets, even worse than using single client data for training for most of the clients.
In ABIDE I, FedBN excell by a non-negligible margin on three clients regarding the mean testing accuracy.
The results are inspiring and bring the hope of deploying FedBN to healthcare ﬁeld, where data are often limited, isolated and heterogeneous on features.
C ONCLUSION AND D ISCUSSION
This work proposes a novel federated learning aggregation method called FedBN that keeps the local
Batch Normalization parameters not synchronized with the global model, such that it mitigates feature shifts in non-IID data.
We provide the convergence guarantees for FedBN in realistic federated settings under overparameterized neural networks regime, while also accounting for practical issues.
In our experiments, our evaluation across a suite of federated datasets has demonstrated that FedBN can signiﬁcantly improve the convergence behavior and model performance of non-IID datasets.
FedBN is independent of the optimization strategy. Thus, it can be used with various optimizers (Reddi et al., 2020). We also note that since FedBN makes only lightweight modiﬁcations to
FedAvg and has much ﬂexibility to be conbimed with other strategies, these merits allow us to easily integrate FedBN into existing tool-kits/systems, such as Pysyft (Ryffel et al., 2018), Google TFF (Google, 2020), Flower (Beutel et al., 2020), dlplatform (Kamp & Adilova, 2020) and FedML (He et al., 2020). We believe that FedBN can improve a wide range of applications.
An interesting direction for future work is analyzing what types of differences in local data can be ironed out by FedBN and where the limits of the method lie. Moreover, not sharing local BN parameters should make inference on local data more challenging. It would be interesting to quantify this improvement in privacy. A N OTATION T ABLE
Notations
Description x features, x ∈ R d d dimension of x y labels, y ∈ R
P ( · ) probability distribution
N total number of clients
T total number of epochs in training
E number of local iteration in FL
M number of training samples in each client [ N ] set of numbers, [ N ] = { 1 , . . .
, N } i indicator for client, i ∈
[ N ] j indicator for sample in each client, j ∈ [ M ] ( x ij , y ij ) the j -th training sample in client i m number of neurons in the ﬁrst layer k indicator for neuron,
[Table]
] v k parameters for the k
[Table]
neuron in the ﬁrst layer
[Table]
, given a matrix S
S i covariance matrix for features in client i , S i
=
E x i x i (cid:62) p, q indicator for sample, p, q ∈ [ NM ] f two layer ReLU neural network with BN f ∗ two layer ReLU neural network with BN with client-speciﬁed BN parameters
V parameters of the ﬁrst phase neurons, V ∈ R m ×
d γ the scaling parameter of BN c top layer parameters of the network σ ( · )
ReLU activation function, σ ( · ) = max {· , 0 } N ( µ , Σ )
Gaussian with mean µ and covariance Σ
U [ − 1 , 1]
Rademacher distribution α variance of v k at initialization
L ( f ) empirical risk with square loss for network f Λ ( t ) evolution dynamic for FedAvg at epoch t
V ( t ) evolution dynamic with respect to V for FedAvg at epoch t
G ( t ) evolution dynamic with respect to γ for FedAvg at epoch t Λ ∗ ( t ) evolution dynamic for FedBN at epoch t
V ∗ ( t ) evolution dynamic with respect to V for FedBN at epoch t
G ∗ ( t ) evolution dynamic with respect to γ for FedBN at epoch t λ min ( A ) the minimal eigenvalue of matrix
A
G ∞ expectation of G ( t )
G ∗∞ expectation of G ∗ ( t )
B C ONVERGENCE P ROOF B. E VOLUTION D YNAMICS
In this section
[Equation]
we calculate the evolution dynamics Λ ( t ) for training with function f and Λ ∗ ( t ) for training with f ∗ .
Since the parameters are updated using gradient descent, the optimization dynamics of parameters are
[Equation]
Let f p = f
( x i p p ) .
Then, the dynamics of the prediction of the p -th data point in site i p is
[Equation]
(cid:124)
[Equation]
The gradients of f p and L with respect to v k and γ k are computed as
[Equation]
where f p = f
( x i p p ) , x v ipk (
[Equation]
) ⊥ p
[Equation] [Equation]
and 1 pk
[Equation]
We deﬁne Gram matrix V ( t ) and G ( t ) as
[Equation]
It follows that
[Equation]
− y q ) .
Let f = ( f 1 , . .
. , f n ) (cid:62) =
( f ( x 1 ) , . . .
, f ( x NM )) (cid:62) .
The full evolution dynamics are given by
[Equation]
where Λ ( t ) :
= V ( t ) α 2 + G ( t ) .
We compute the NTK for FedBN with f ∗ analogously.
We deﬁne the Gram matrix V ∗ ( t ) and G ∗ ( t ) as
[Equation]
The full evolution dynamics are then
[Equation]
where Λ ∗ ( t ) :
= V ∗ ( t ) α+ G ∗ ( t ) .
B. P ROOF OF L EMMA 4.
Dukler et al.
(2020) proved that the matrix G ∞ is strictly positive deﬁnite.
In their proof, G ∞ is the covariance matrix of the functionals φ p deﬁne as
[Equation]
over the Hilbert space V of L 2 (cid:0) N (cid:0) 0 , α 2 I (cid:1)(cid:1) .
G ∗∞ is strictly positive deﬁnite by showing that φ 1 , · · · , φ NM are linearly independent, which is equivalent to that
[Equation]
holds only for c p = 0 for all p .
Let G ∞ i be the i -th M × M block matrices on the diagonal of G ∞ .
Then G ∗∞ = diag ( G ∞ 1 , · · · , G ∞ N ) .
To prove that G ∗∞ is strictly positive deﬁnite, we will show that G ∞ i is positive deﬁnite.
Deﬁne
[Equation]
Then, we are going to show that
[Equation]
(11) holds only for c j = 0 for all j ∈
[ M ] .
Suppose there exist c 1 , · · · , c M that are not identically 0, satisfying (11).
Let the coefﬁcients for client i c be the same as c 1 , · · · , c M and let the coefﬁcients for other client be 0.
Then, we have s sequence of coefﬁcients satisfying (10), which is a contradiction of that G ∞ is strictly positive deﬁnite.
This implies G ∞ i is strictly positive deﬁnite, which means the eigenvalues are positive.
Since the eigenvalues of G ∗∞ are exactly the eigenvalues of G ∞ i , λ min ( G ∗∞ ) is positive and thus, G ∗∞ is strictly positive deﬁnite.
B. P ROOF OF C OROLLARY 4.
To compare the convergence rates of FedAvg and FedBN when E = 1 , we compare the exponential factor in the convergence rates, which are (1 − ηµ 0 / 2) and (1 − ηµ ∗ 0 / 2) for FedAvg and FedBN, respectively.
Then, it reduces to comparing µ 0 =
λ min ( G ∞ ) and µ ∗ 0 = λ min ( G ∗∞ ) .
Comparing equation (7) and (9), G ∗∞ takes the M × M block matrices on the diagonal of G ∞ :
[Equation]
where G ∞ i is the i -th M × M block matrices on the diagonal of G ∞ .
By linear algebra,
[Equation]
Since the eigenvalues of G ∗∞ are exactly the union of eigenvalues of G ∞
i , we have
[Equation]
Thus, (1 − ηµ 0 / 2)
C F ED BN A LGORITHM
We describe the details algorithm of our proposed FedBN as following Algorithm 1:
Notations: The user indexed by k , neural network layer indexed by l , initialized model parameters: w ( l ) 0 ,k , local update pace: E , and total optimization round T .
1: for each round t = 1 , 2 , . .
, T do 2: for each user k
and each layer l do 3: w ( l ) t +1 ,k ← SGD ( w ( l ) t,k ) 4: end for 5:
if mod ( t, E ) = 0 then 6: for each user k
and each layer l do 7: if layer l is not BatchNorm then 8: w
[Table]
9: end if 10: end for 11: end if 12: end for
D E XPERIMENTAL DETAILS D. V ISUALIZATION OF B ENCHMARK D ATASETS
[Figure]
D. M ODEL A RCHITECTURE AND T RAINING D ETAILS ON B ENCHMARK
We illustrate our model architecture and training details of the digits classiﬁcation experiments in this section.
Model Architecture.
For our benchmark experiment, we use a six-layer Convolutional Neural
Network (CNN) and its details are listed in Table 3.
Layer
Details
[Table]
BN
[Table]
MaxPool2D
[Table] [Table]
BN
[Table]
Training Details.
We give detailed settings for the experiments conducted in 5.1: (1) convergence rate (Table 4), (2) analysis of local update epochs (Table 5), (3) analysis of local dataset size (Table 6), (4) effects of statistical heterogeneity (Table 7) and (5) comparison with state-of-the-art (Table 8).
Each table describes the number of clients, samples and the local update epochs.
During training process, we use SGD optimizer with learning rate 10 − 2 and cross-entropy loss, we set batch size to 32 and training epochs to 300.
For hyper-parameter µ , we use the best value µ = 10 − 2 founded by grid search from the the default settings in FedProx Li et al.
(2020b).
[Table]
[Table] [Table]
The
[Table]
of samples ω ∈ { 74 , 371 , 743 , 1487 ,
[Table] [Table]
the range from 1 to 10.
We increase number of clients step by step and number of samples will increase accordingly
[Table]
D.
M ODEL A RCHITECTURE AND T RANING D ETAILS OF I MAGE C LASSIFICATION T
ASK
ON O FFICE -C ALTECH 10 AND D OMAIN N ET
In this section, we provide the details of our model and training process on both Ofﬁce-Caltech
Gong et al.
and DomainNet Peng et al.
(2019) dataset.
Model Architecture.
For the image classiﬁcation tasks on these two real-worlds datasets OfﬁceCaltech
and DomainNet data, we use adapted AlexNet added with BN layer after each convolutional layer and fully-connected layer (except the last layer), architecture is shown in Table
[Table]
fully connected layer (FC), we list input and output dimension.
For BatchNormalization layer (BN), we list the channel dimension.
Training Details.
Ofﬁce-Caltech10 selects 10 common objects in Ofﬁce-31 Saenko et al.
and Caltech-256 datasets Grifﬁn et al.
There are four different data sources, one from
Caltech-256 and three from Ofﬁce-31, namely Amazon(images collected from online shopping website), DSLR and Webcam(images captured in ofﬁce environment using Digital SLR camera and web camera).
We ﬁrst reshape input images in the two dataset into 256 × 256 × 3, then for training process, we use cross-entropy loss and SGD optimizer with learning rate of 10 − 2 , batch size is set to 32 and training epochs is 300.
When comparing with FedProx, we set µ to 10 − 2 which is tuned from the default settings.
The data sample number are kept into the same size according to the smallest dataset, i.e. Ofﬁce-Caltech10 uses 62 training samples and DomainNet uses 105 training samples on each dataset.
In addition, for simplicity, we choose top-10 class based on data amount from DomainNet containing images over 345 categories.
D. ABIDE D ATASET AND T RAINING D ETAILS
Here we describe the real-world medical datasets, the preprocessing and training details.
Dataset:
The study was carried out using resting-state fMRI (rs-fMRI) data from the Autism
Brain Imaging Data Exchange dataset (ABIDE I preprocessed, (Di Martino et al., 2014)).
ABIDE is a consortium that provides preciously collected rs-fMRI ASD and matched controls data for the purpose of data sharing in the scientiﬁc community.
We downloaded Regions of Interests (ROIs) fMRI series of the top four largest sites (UM, NYU, USM, UCLA viewed as clients) from the preprocessed ABIDE dataset with Conﬁgurable Pipeline for the Analysis of Connectomes (CPAC) and parcellated by Harvard-Oxford (HO) atlas.
Skipping subjects lacking ﬁlename, resulting in 88, 167, 52, 63 subjects for UM, NYU, USM, UCLA separately.
Due to a lack of sufﬁcient data, we used sliding windows (with window size 32 and stride 1) to truncate raw time sequences of fMRI.
The compositions of four sites were shown in Table 10.
The number of overlapping truncate is the dataset size in a client.
[Table]
Training Process :
For all the strategies, we set batch size as 100.
The total training local epoch is 50 with learning rate 10 − 2 with SGD optimizer.
Local update epoch for each client is E = 1 .
We selected the best parameters µ = 0 .
2 in FedProx through grid search.
E M ORE E XPERIMENTAL R ESULTS ON B ENCHMARK D ATASETS E. C ONVERGENCE C OMPARISON OVER F ED A VG AND F ED BN
In this section we conduct an additional convergence analysis experiment over different local update epochs settings: E = 1 , 4 , 8 , 16 .
As shown in Fig.
7, FedBN converges faster than FedAvg under different values of E , which is supportive to our theoretical analysis in section 4 and experimental results in section 5.
E. D ETAILED S TATISTICS OF F IGURE
[Table]
E. C OMPARE F ED BN WITH C ENTRALIZED T RAINING
To better understand the signiﬁcance of the numbers reported in our main context, we compare
FedBN with centralized training, that pools all training data in to a center.
We present the testing accuracy on each digit dataset in Table 12.
FedBN, federated learning with data-speciﬁc BN layers, even performs better than vanilla centralized training strategy.
[Table]
E. D IFFERENT C OMBINATIONS OF E AND B
In this section, we show different combinations of local update epochs E and batch size B .
Specifically, E ∈ { 1 , 4 , 16 } and B ∈ { 10 , 50 , ∞} , ∞ denotes full batch learning.
Following the setting in original FedAvg paper McMahan et al.
(2017), we present the comparisons between FedBN and
FedAvg on each combination of E and B in Table 13.
The results are in good agreement that FedBN can consistently outperform FedAvg and robust to batch size selection.
Further, we depicts the test sets accuracy vs. communication epochs under different combination of E and B in Figure 8.
In conformity with the observation in McMahan et al.
(2017), increasing local update frequency by decreasing E can produce a striking decrease on the communication round to meet a given accuracy target (e.g., the gray line in Figure 8).
B does not have a clear trend of affecting communication cost in Figure 8, but smaller B achieved better ﬁnal accuracy as shown in Table 13.
[Table]
MNIST
B=10, E=
FedAvg 72.
FedBN 78.
B=10, E=
FedAvg 75.
FedBN 78.
B=10, E=
FedAvg 73.
FedBN 79.
B=50, E=
FedAvg 69.
FedBN 75.
B=50, E=
FedAvg 71.
FedBN 76.
B=50, E=
FedAvg 70.
FedBN 76.
B= ∞ , E=
FedAvg 68.
FedBN 73.
B= ∞ , E=
FedAvg 71.
FedBN 76.
B= ∞ , E=
FedAvg 71.
FedBN 74.
E. D ETAILED S TATISTICS OF V ARYING L OCAL D ATASET S IZE E XPERIMENT
Considering putting all results in one ﬁgure (15 lines) might affect readability of the ﬁgure, we excluded the statistics of FedAvg in our Fig.
4 (b), the ablation study of our method on the effect of local dataset size.
Here, we list the full results in Table 14.
It is not too surprising that at Singleset can be the best when the a local client gets a lot of data.
Setting 100% 60% 40% 20% 10% 5% 1%
MNIST
SingleSet 98.
FedAvg 98.
FedBN 98.
SVHN
SingleSet 86.
FedAvg 82.
FedBN 87.
USPS
SingleSet 98.
FedAvg 98.
FedBN 99.
Synth
SingleSet 94.
[Table]
E.
T RAINING ON U NEQUAL D ATASET S IZE
In our benchmark experiment (Section 5.1), we truncate the sample size of the ﬁve datasets to their smallest number.
This data preprocessing intends to strictly control non-related factors (e.g., imbalanced sample numbers across clients), so that the experimental ﬁndings can more clearly reﬂect the effect of local BN.
In this regard, truncating datasets is a reasonable way to make each client have an equal number of data points and local update steps.
It is also possible to keep the data sets in their original size (which is unequal), by allowing clients with less data to repeat sampling.
In this way, all clients use the same batch size and same local iterations of each epoch.
We add results of such a setting with 10% and full original datasize in Table 15 and Table 16 respectively.
It is observed that
FedBN still consistently outperforms other methods.
[Table] [Table]
F S YNTHETIC D ATA E XPERIMENT
Settings We generate data from two-pair of multi-Gaussian distributions.
For one pair, samples ( x, 0) and ( x, 1) are sampled from N ( − 1 , Σ 1 ) and N (1 , Σ 1 ) respectively, with coveriance Σ 1 ∈ R× 10 .
For another pair, samples ( (cid:101) x, 0) and ( (cid:101) x, 1) are sampled from N ( − 1 , Σ 2 ) and N (1 , Σ 2 ) respectively, with coveriance Σ 2 ∈ R 10 × 10 .
Speciﬁcally, we design convariance matrix Σ 1 as an identity diagonal matrix and Σ 2 is different from Σ 1 by having non-zero values on off-diagonal entries.
We train a two-layer neural network with 100 hidden neurons for 600 steps using crossentropy loss and SGD optimizer with 1 × 10 − 5 learning rate.
Denote W k and b k are the in-connection weigths and bias term of neuron k .
We initialize the model parameters with W k ∼ N (0 , α 2 I ) , b k ∼
N (0 , α 2 ) , where α = 10 .
Results.
The aim of synthetic experiments is to study the behavior of using FedBN with a controlled setup.
We achieve 100% accuracy on binary classiﬁcation for FedAvg and FedBN.
Fig.
9 shows comparison of training loss curve over steps using FedAvg and FedBN, presenting that FedBN obtains signiﬁcantly faster convergence than FedAvg.
[Figure]
G T RANSFER L EARNING AND T ESTING ON U NKNOWN D OMAIN C LIENT
In this section, we discuss out-of-domain generalization of FedBN and prove the solutions for the following two scenarios:
1) transferring FedBN to a new unknown domain clients during training; 2) testing an unknown domain client.
If a new center from another domain joins training, we can transfer the non-BN layer parameters of the global model to this new center.
This new center will compute its own mean and variance statistics, and learn the corresponding local BN parameters.
[Table]
Testing the global model on a new client with unknown statistics outside federation requires allowing access to local BN parameters at testing time (though BN layers are not aggregated at the global server during training).
In this way, the new client can use the averaged trainable BN parameters learned at existing FL clients, and compute the (mean, variance) on its own data.
Such a solution is also in line with what was done in recent literature, e.g., SiloBN (Andreux et al., 2020).
We conduct the experiment with this solution for FedBN and compared its performance with FedAvg and FedProx.
Speciﬁcally, we use the digits classiﬁcation task and treat the two unseen datasets –
Morpho-global and Morpho-local from Morpho-MNIST (Castro et al., 2019) as the two new clients.
The new clients contain substantially perturbed digits.
Speciﬁcally, Morpho-global containing thinning and thickening versions of MNIST digits, while Morpho-local changes MNIST by swelling and fractures.
The results are listed in Table 17.
It is observed that the obtained results from three methods are generally comparable in such a challenging setting, with FedBN presenting slightly higher performance on overall average accuracy.