FedBN: Federated Learning on Non-IID Features via Local Batch Normalization
Authors: Xiaoxiao Li
Meirui JIANG
Xiaofei Zhang
Michael Kamp
Qi Dou
Venue: ICLR 2021 Poster (2021)
Keywords: Federated Learning
Non-IID
Batch Normalization

Abstract: 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 significantly 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 shifts. Unlike those settings, we address an important problem of FL, e.g., different scanners/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 simplified setting that FedBN has a faster convergence rate than FedAvg. Code is available at https://github.com/med-air/FedBN.

Published as a conference paper at ICLR 2021
FEDBN: FEDERATED LEARNING ON NON-IID FEATURES VIA LOCAL BATCH NORMALIZATION
Xiaoxiao Li
Meirui Jiang
Department of Computer Science Princcton University x xiaoxias ac 11 Baya. yale.
Department of Computer Science and Engineering The Chinese University of Hong Kong mr j ange cse. www.cuckkkoooooo .ddu.
Xiaofei Zhang
Michael Kamp
Department of Statistics Iowa State University
Dept of Data Science and AI, Faculty of IT Monash University
xf #fzhang@isstate.cc edu
michael. kamp edu
Qi Dou
Department of Computer Science and Engineering The Chinese University of Hong Kong qdouêcse. cuhk. edu .hk
ABSTRACT
The emerging paradigm of federated learning (FL) strives enable collabora- tive training of deep models on the network edge without centrally aggregating data and hence improving data privacy. In most cases, the assumption of in- dependent and identically distributed samples across local clients does not hold for federated learning setups. Under this setting, neural network training per- formance may vary significantly according to the data distribution and even hurt training convergence. Most of the previous work has focused on difference the distribution of labels or client shifts. Unlike those settings, we address an important problem of FI., e.g., different scanners/sensors in medical imaging, dif- ferent scenery distribution in autonomous driving (highway city), where local clients store examples with different distributions compared to other clients, which denote feature shift non-iid. In this work, propose 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-arr for non-iid data (FedProx) on our extensive experi- ments. These empirical results supported by convergence analysis that shows in a simplified setting that FedBN has a faster convergence rate than FedAvg, Code is available https://github.coo. wttp://ooooommmmiir/oddooooo
INTRODUCTION
Federated learning (FL), has gained popularity for various applications involving learning from dis- tributed data. In FL, 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 al., 2019; et al., 2020b). 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., 201 8; 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.

---

Published as a conference paper at ICLR 2021
T-
16-
120
7- dast Iwth BX 2:15
1.90
1.11
6.00
0.51
0.5
01-
Sent
0:40
soodd
Figure Training error on local datasets for two clients respectively with and w/o BN, where BN harmonizes the loss surface.
Figure 2: Error surface of a client for model parameter E [0.001, 12] and BN parameter 7 E [0.001, 4]. Averaging model and BN pa- rameters leads to worse solutions.
in the feature space, which has not yet been explored in the literature. Specifically, consider that local data deviates in terms of the distribution in feature space, and identify this scenario as feature shifi. This type of non-iid data 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 different imaging machines and protocols used in hospitals, e.g., different intensity and contrast. In this example, each hospital is 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 practice is utilizing Batch Normalization (BN) (loffe & Szegedy, 2015): recent work has proposed BN 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 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 FL. Toy Example: We consider simple non-convex learning problem: we generate data v. y E IR with = cos( (W)rue + where E IR is drawn iid from Gaussian distribution and is zero-mean Gaussian noise and consider models of the form f (x) = cos (wx)) with model parameter local batch normalization harmonizes local data distributions. We consider a simplified form of BN that normalizes the input by scaling it with y, i..., the local empirical standard deviation, and a setting with 2 clients. As Fig. shows, the local squared loss is very different between the two clients. Thus, averaging the model does not lead to 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 both model parameter E IR and BN parameter 7 E IR in Fig. 2. The figure shows that for an optimal weight wi. changing 7 deteriorates the model quality. Similarly, for a given optimal BN parameter 77 changing w deteriorates the quality. In particular, the average model w (wi + w / and average BN parameters ア = (>i + 752 has high generalization error. At the same time, the average model w with local BN parameter 7i performs very well.
R. Local data deviates in the variance of æ. First, we illustrate that
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 up- dated 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 aggregation scheme. This approach has zero parameters to tune, requires minimal additional computational re- sources, 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 pver-parameterized regime. In addition,
2

---

Published as a conference paper at ICLR 2021
have conducted extensive experiments benchmark and three real-world datasets. Compared to classical FedAvg, as well as the state-of-the-ar 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 heterogeneous over local client. Empirical work addressing non-iid issues, mainly focus 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), recent framework tackled the heterogeneity by allowing partial information aggregation and adding a proximal term to FedAvg. Zhao et al. (2018) assumed subset of the data is globally shared between all the clients, hence generalizes the problem hand. FedMA (Wang et al., 2020) proposed an aggregation strategy for non-iid data partition that shares global model in 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 affine distribution shift and tackles this problem by learning the affine transformation. This hampers the generalization when we cannot 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 al., 2020; Zhang et al., 2020).
Batch Normalization in Deep Neural Networks: Batch Normalization (loffe & 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 normal- ization. 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 sin- gle direction reliance, thus improving model generalizability. Li et al. (2018) took advantage of BN for tackling the domain adaptation problem. However, what role BN is playing in the scope of federated learning, especially for non-iid training, still remains unexplored to date.
PRELIMINARY
Non-IID Data in Federated Learning: We introduce the concept of feature shift in federated learning 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 (x y) P((y))) and P. (x(x)) P. We define feature shift as the case that covers: 1) covariate shift: the marginal distributions (x( varies across clients, even if (r x) is the same for all client; and 2) concept shift: the conditional distribution (. varies across clients and Pyy) 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 casiest to implement federated learning strategy, where clients collaboratively send updates of locally trained models to global server. Each client runs local copy of the global model on its local data. The global model's '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 shared prediction model while keeping all training data locally. Although FedAvg g 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).
3

---

Published as a conference paper at ICLR 2021
4 FEDERATED AVERAGING WITH LOCAL BATCH NORMALIZATION
4.1 PROPOSED METHOD FEDBN
We propose 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 simplified scenario, which we show that FedBN improves the convergence rate under feature shift.
4.2 PROBLEM SETUP
We assume N E 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 E N training examples (a difference in training examples can be account for by weighted averaging (MeMahan al., 2017)) for regression task, i.e., each client i E [N] ([N]= (1,..., N)) has training examples [(x), vt)) E IRd xII j E [M]] Furthermore, we assume two-layer neural network with ReLU activations trained by gradient descent. Let væ Rd denote the parameters of the first layer, where & E [m] and m is the width of the hidden layer. Let
v lls vv 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 E [N] the inputs x3 are centered (Ex' 0) with covariance matrix S, Ex'x'T where S; independent from the label y and may differ for each i E [N] e.g., S; are not all identity matrices, and for each index pair # Xx & R XxXgJe for all [0].
With Assumption 4.1, the normalization of the first layer for client i v.x'' IIve lls; FedBN with client-
specified BN parameters trains model f Rd + parameterized by (V, e ma x pmmNN, x Rmm i.e.
m Ck
V
) 1 1(x client
T
Vk lls.
where y 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 FedAvg instead trains function f: Rd + I which is a special case of Eq. with YE,i T' for Vi E [N]. We take a random initialization of the parameters (Salimans Kingma, 2016) in our analysis:
(0) ~ N (0, Ck ~ U[-1,,, and 7hk = Th.; = v/ (0)l2/
(2)
where a² controls the magnitude of v4 at initialization. The initialization of the BN parameters Yk and y independent of The parameters of the network fx; V, e) are obtained by minimizing the empirical risk with respect to the squared loss using gradient descent
M
L( NM
(M(x))-
(3)
=1
4.3 CONVERGENCE ANALYSIS
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. (2018). Recent machine learning theory
We omit centering neurons as well learning the shift parameter of BN for the neural network analysis because of the assumption that is zero-mean and the two layer network setting (Kohler al., 2019; Salimans Kingma, 2016).
4

---

Published as a conference paper at ICLR 2021
studies (Arora al., 2019; Du et al., 2018; Allen-Zhu et al., 2019; van den Brand al., 2020; Dukler et al., 2020) have shown that for finite-width perr-parmtterized 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 magnitude component (t) and direction component V(t)/a² following Dukler et al. (2020):
df di (()((())) - y), where A()) = V(t) a²
+ G(I).
The specific forms of (t) and G(t) are given in Appendix B.1. Let Amin (A) denote the minimal eigenvalue of matrix A. The matrices V(t) and G(t) are positive semi-definite, since they can be viewed as covariance matrices. This gives Xmin (A(t)) Z max {Amin (V(t))/a² Amiu (G(t)) } According to NTK, the convergence rate is controlled by Amin Amin.((()) (A((). Then, for a >> 1, convergence dominated by (t). Let (1) and A (t) denote the evolution dynamics of FedAvg and FedBN and let Gt) and G (t) denote the magnitude component in the evolution dynamies of FedAvg and FedBN, For the convergence analysis, we use the auxiliary version of the Gram matrices, which is defined as follows.
NNM we define the auxiliary Gram matrices G 0o Definition 4.2. Given sample points {xp)== BNMMN and G E BNMNNN as
G = E (v xp) (v xxxx) (FedAvg)
(4)
pq G := Bv-N(0, (v T xp) (v x.y 1ipp ig), (FedBN).
(5)
Given Assumption 4.1, we use the key results Dukler et al. (2020) to show that G is positive definite. Further, we show that to is positive definite. We use the fact that the distance between G(t) and its auxiliary version is small in over wver-paaametrizeed neural network, such that Gtt) remains positive definite.
Lemma 4.3. Fix points {xp)p=1
NNM satisfying Assumption 4. Then Gram matrices G and G defined as (4) and (5) are strictly positive definite. Let the least eigenvalues be Amin A =: μoo
and Amaa (G =: 10: where Po: P0 O.
Proof sketch The main idea follows Du et al. (2018); Dukler et al. (2020), that given points {xp ) N.M the matrices G and T can be shown covariance matrix of linearly independent operators. More details 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. (2020) 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. (2020)). Suppose network (4) initialized as in (2) with a > 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 I1 = 0(1). If m = Ω (max ( (N*M log(N (/8)NM)))/ pby N2 M log(N (////1)) (.)). probability
then with
1. For iterations t = 0,1, the evolution matrix Att) satisfies Amiu (A()))())))))) 2
Re. 2
2. Training with gradient descent of step-size 1) O
converges linearly
|
[|E(t) - -112 <
171000
||F(O)= -1.
Following the key ideas in Dukler et al. (2020), here we further characterize the convergence for FedBN.
Corollary 4.5 (G-dominated convergence for FedBN), Suppose network (5) and all other condi- tions in Theorem 4.4. With probability 1- 8, for iterations 0, 1, the evolution matrix
(1) satisfies Amio ( A (()) Z Aé and training with gradient descent of step-size 1) O [T-(m]]]
converges linearly as (1) - -1 < (1 (1- m4u [(((0) - -y00--ll: 11.
5

---

Published as a conference paper at ICLR 2021
SVHN
USPS
SymDOggt
MNIST-M
MMSST
FodAug
Fodeug FediN
Foddug FedeN
FedAvg FedUN
Fee/vg
2.0
2.3
>0
20
1.0
100
1.0
305
200
102
400
300
600
100
200
106
700
306
200
400
130
200
400
500
120
300
300
400
500
138
Figure 3: Convergence of the training loss of FedBN and FedAvg on the digits classification datasets. FedBN exhibits faster and more robust convergence.
The exponential factor of convergence for FedAvg (1 ()110/2 and FedBN (1 - 725/2) are con- trolled by the smallest eigenvalue of G(), respectively G* Then we can analyze the convergence performance of FedAvg and FedBN by comparing Amin (G 00) and m
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 Amin( (G Amin (G*0e) Comparing equation (4) and (5), G takes the M x M block matrices on the diagonal of o Let G be the i-th M X M block matrices on the diagonal of - By linear algebra, Amin ( G 2 (G for i E [N]. Since G -o = diag( G the result Amin (G 2 Amin (Go).
G we have Amin (G = [ Amin )) Therefore, we have
5 EXPERIMENTS
In this section, we demonstrate that using local BN parameters beneficial 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.1 BENCHMARK EXPERIMENTS
Settings: We perform an extensive empirical analysis using a benchmark digits classification 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 dis- tribution. Specifically. we use the following five datasets: SVHN Netzer et al. (2011), USPS Hull (1994), SynthDigits Ganin & Lempitsky (2015), MNIST-M Ganin & Lempitsky (2015) and MNIST LeCun et al. (1998). To match the setup in Section 4, we truncate the sample size of the five datasets their smallest number with random sampling, resulting in 7438 training samples in each dataset 2 Testing samples are held out and kept the same for all the experiments on this benchmark dataset. Our classification model is convolutional neural network where BN layers are added following each feature extraction layer (i.e., both convolutional and fully-connected). The architecture is de- tailed 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 specified, 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, the FI. system contains five clients. Each client exclusively owns data sampled from one of the five datasets. More details are listed in Appendix D.2.
"This data preprocessing intends to strictly control non related factors (e.g.. imbalanced sample numbers across clients), so that the experimental findings can more clearly reflect the effect of local BN, Results without truncating reported in Appendix E.2.
'Choosing 10% fraction as default is based on (i) considering it as a typical setting to present general efficacy of our method: (ii) matching literature where the client size is usually around 100 to 1000 data points (McMahan et al., 2017; Li et al., 2019; Hsu et al., 2019), which is similar scale with respect our 10% setting (in terms of the absolute value of sample numhers).

---

Published as a conference paper at ICLR 2021
0.9
100
100
0.8
-9
70
securac
0.7
Accuracy)
AccuracylN::
0.54
0.6
0.82
40
5MH
0.5
Lmg
Emg
E+1
Nedtig NediNN
20
Average
SynchDigtt MNIST
E-16
10
60
0.1
150
20
Epochs
(b)
(0)
Figure Analytical experimental results on: (a) Analysis on different local updating epochs.
FedBN consistently outperforms FedAvg in testing accuracy. (b) Model performance over vary- ing dataset size on local clients. (c) Testing accuracy on different levels hetcrogeneity.
the following paragraphs, we present a comprehensive investigation on the prop- erties of the proposed FedBN approach, including: (I) 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.c., training an individual model within each client). In Appendix G, also provide empirical results for including new client with data from an unknown domain into the learning system.
Overviews:
Convergence Rate: We analyze the training loss curve of FedBN in comparison with FedAvg, as shown in Fig. 3. 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 more comprehensive comparison with dif- ferent local update epochs F. on convergence rate of FedBN and FedAvg (see Appendix F, The results show similar patterns as in Fig 3.
Aggregating 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 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 final testing accuracy, FedBN's accuracy stably exceeds the accuracy of FedAvg on various E.
Analysis of Local Updating Epochs:
1, 4, 8, 16
Analysis 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 significantly 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 benefit from collaborative training on distributed data, especially when each client only holds a small amount of data which are non-iid.
Effeets Statistical Heterogeneity: salient question that arises is: to what degree of het- erogeneity feature shift FedBN is superior to FedAvg. To answer the question, we simulate federated settings with varying heterogeneity as described below. We parcel each dataset into 10 subsets, 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 non-iid clients.
*Dtailled statistics and FedAvg results are presented in Appendix F.2.

---

Published as a conference paper at ICLR 2021
Callech-
Domainnel
ABIDE (medical)
Method
A
D
W
C
P
Q
R
S
NYU
USM
UM
UCLA
SingleSet
54.9
40.2
78.7
86.4
41.0
23.8
36,2
73.1
48.5
34.0
58.0
73.4
64.3
57.3
(1.5)
(1.6)
(1.3)
(2.4)
(0.9)
(1.2)
(2.7)
(0.9)
(1.9)
(1.1)
(3.3)
(2.2)
(1.4)
(2.4)
FedAvg
54.1
44.8
66.9
85.1
48.8
24.9
36.5
56.1
46.3
36.6
62.7
73.1
70.7
64.7
(1.1) (1.0)
(1.5)
(1.9
(0.7)
(1.1)
(1.6)
(1.4)
(2.5)
(2.4)
(0.7)
54.2 44.5
65.0
84.4
48.9
24.9
36.6
54.4
47.8
36.9
63.3
73.0
70.5
64.5
FedProx
(2.5) (0.5)
(3.6)
(1.7)
(0.8)
(1.0)
(1.8)
(3.1)
(0.8)
(2.1)
(1.0)
(1.8)
(1.1)
(1.2)
FedBN
63.0
45.3
83.1
90.5
5122
26.8
41.5
71.3
54.8
42.1
65.6
75.1
68.6
65.5
(1.6) (1.5)
(2.5)
(2.3)
(1.4)
(0.5)
(1.4)
(0.7)
(0.8)
(1.3)
(1.1)
(1.4)
(2.9)
(1.0)
Table 1: We report results on three different real-world datasets with format mean(std) from 5-trial
rum. For Office-Caltech Caltech 10, A, C, D, are abbreviations for Amazon, Caltech, DSLR and WebCam, for DomainNet, C. I, P. Q, R, S are abbreviations for Clipart, Infograph, Painting, Quickdraw, Real and Sketch. For ABIDE, we list the abbreviations for the clients (i.e., medical institutions).
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 E {1,.... 9}$ 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 accuracy than FedAvg over all levels of heterogeneity.
Comparison with State-of-the- art: 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 benefit of easy adaptation current FL frameworks in practice. We also include training Sin- gleSet 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 each dataset over trials shown in Fig. 5 From the results, following observation: (1) FedBN achieves the highest accuracy, consistently outperforming the state-of-the-art and baseline methods; (2) FedBN achieves the most significant improvements SVHN whose image appearance is very different from others (i.c., presenting more obvious feature shift); (3) FedBN shows a smaller variance in error over multiple runs, indicating its stability.
100
SingleSet FedAng FediN
Accuracy%)
60
SVIIN
USPS
SynthDiigits Datasets
MNIST-M
MNIST
Figure 5: Performance on benchmark experiments
make the
5.2 EXPERIMENTS ON REAL-WORLD DATASETS
To better understand how our proposed algorithm can be beneficial in real-word feature-shit non- iid, we have extensively validated the effectiveness of FedBN in comparison with other methods on three real-world datasets: image classification on 1ffice-Cattec10 (Gong et al., 2012) with images acquired in different cameras or environments; image classification on DomainNet (Peng et al., 2019) with different image styles; and a neurodisorder diagnosis task on ABIDE (Di Martino et al., 2014) with patients from different medical institutions".
NNaeely, each increment contains five non-iid clients.
Detailed statisties are shown in Appendix E.2 Mooe details about the datasets and training process listed in the Appendix D.3 and D.4
8

---

Published as a conference paper at ICLR 2021
Datasets and Setup: (1) We conduct the classification task natural images from Office- Caltech10, which has four data sources composing Office-31 Saenko et al. (2010) (three data sources) and Caltech-256 datasets (one data source) Griffin et al. (2007), which are acquired us- ing 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 (I), 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 that collects functional brain images using differ- ent imaging equipment and protocols. We validate on medical application for binary classification between autism spectrum disorders patients and healthy control subjects.
The
10 contains ten categories of objects. The DomainNet extensively contains 345 object categories and we use the top ten most common classes to form sub-dataset for our ex- periments. Our classification models adopt AlexNet (Krizhevsky et al., 2012) architecture with BN added after each convolution and fully-connecter layer. Before feeding into the network, all images are resized to 256 x 256 x 3. For ABIDE I, each instance is represented as a 5995- 9995-dimnssiona vector through brain connectome computation. We use three-layer fully connected neural network as the classifier with the hidden layers of 16 with two BN layers after the first two fully connected layers. Same as the above benchmark, we perform 5 repeated runs for each experiment.
The experimental results are shown in Table 1 in the form of mean (std). On Office-Caltech FedBN significantly outperforms the state-of-the-ar 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 find 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 field, where data are often limited, isolated and heterogeneous on features.
Results and Analysis:
6 CONCLUSION AND DISCUSSION
This work proposes 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-lll) data. We provide convergence guarantees for FedBN in realistic federated settings under the overparameterized neural networks regime, while also accounting for practical issues. In our experiments, that FedBN can significantly improve the convergence behavior and model performance of non-lID datasets. We also demonstrate the effectiveness of FedBN in scenarios that where a new client with an unknown domain joins the FL system (see Appendix G).
evaluation across suite of federated datasets has demonstrated
FedBN is independent of the communication and aggregation strategy and thus in practice be readily combined with different optimization algorithms, communication schemes, and aggregation techniques. The theoretical analysis of such combinations is interesting direction for future work. We also note that since FedBN makes only lightweight modifications to FedAvg and has much flexibility be combined with other strategies, these merits allow us easily integrate FedBN into existing tool-kits/systems, such as Pysyft (Ryffel et al., 2018), Google TFF (Google, 2020), Flower (Beutel et al., 2020), diplatform (Kamp & Adilova, 2020) and FedML (He et al., 2020)*.
We believe that FedBN can improve a wide range of applications such as healthcare (Ricke et al., 2020) and autonomous driving (Kamp et al., 2018). A few interesting directions for future work include analyzing what types of differences in local data can benefit from FedBN and explore the limits of FedBN. Moreover, privacy is an essential concern in FL. Invisible BN parameters in FedBN should make attacks on local data more challenging. It would be interesting to quantify the privacy- preservation improvement in FedBN.
'The implementations on diplatform and Flower are available, and the implementation on FedMI. is to appear soon.

---

Published as a conference paper at ICLR 2021
REFERENCES
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. International Conference on Machine Learning, pp. 242 252. PMLR, 2019.
Mathieu Andreux, Jean Ogier du Terrail, Constance Beguier, and Eric W Tramel. Siloed federated learning for multi-centric histopathology datasets, In Domain Adaptation and Representation Transfer, and Distributed and Collaborative Learning, pp. 129-139. Springer, 2020.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang- Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. arXiv preprint arXiv:1901.08584 2019.
Daniel Beutel, Taner Topal, Akhil Mathur, Xinchi Qiu, Titouan Parcollet, and Nicholas D Lane. Flower: A friendly federated learning research framework. arXiv preprint arXiv:2007.14390 2020.
Daniel C Castro, Jeremy Tan, Bernhard Kainz, Ender Konukoglu, and Ben Glocker. Morpho-mnist: Quantitative assessment and diagnostics for representation learning. Journal of Machine Learning Research, 20(178):1-29 2019.
Woong-Gi Chang. Tackgeun You, Seonguk Sco, Suha Kwak, and Bohyung Han. Domain-specific batch normalization for unsupervised domain adaptation. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 7354-7362, 2019.
Adriana Di Martino, Chao-Gan Yan, Qingyang Li, Erin Denio, Francisco X Castellanos, Kaat Alaerts, Jeffrey Anderson, Michal Assaf, Susan Bookheimer, Mirella Dapretto, et al. The autism brain imaging data exchange: towards large-scale evaluation of the intrinsic brain archi- tecture autism. Molecular psychiatry, 19(6):659, 2014.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. arXiv preprint arXiv:1810.02054 2018
Yonatan Dukler, Quanquan Gu, and Guido Montúfar. Optimization theory for relu neural networks
trained with normalization layers, 2020.
Yaroslav Ganin and Victor Lempitsky. Unsupervised domain adaptation by backpropagation. In International conference on machine learning. pp. 1 1180-10 1809-189. PMLR, 2015.
Boqing Gong, Yuan Shi, Fei Sha, and Kristen Grauman. Geodesic flow kernel for unsupervised domain adaptation. In 2012 IEEE Conference on Computer Vision and Pattern Recognition, pp. 2066-2073. IEEE, 2012.
Google. TensorFlow Federated: Machine Learning Decentralized Data, 2020. https:t////// tensorflow org/ federated.
Gregory Griffin, Alex Holub, and Pietro Perona. Caltech-256 object category dataset, 2007.
Chaoyang He, Songze Li, Jinhyun So, Zhang, Hongyi Wang, Xiaoyang Wang, Praneeth Vepakomma, Abhishek Singh, Hang Qiu. Li Shen, et al. Fedml: A research library and bench- mark for federated machine learning. arXiv preprint wrXiv22007.13518 2020.
Kevin Hsieh, Amar Phanishayee, Onur Mutlu, and Phillip B Gibbons. The non-iid data quagmire of decentralized machine learning. arXiv preprint 1r0.v0990.00// 2019.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arXiv:1909.0633. 2019.
Jonathan J. Hull. A database for handwritten text recognition research. IEEE Transactions on
pattern analysis and machine intelligence, 16(5):550-5544 1994.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv: arXiv:1502.0033 1502.03167, 2015.
10

---

Published as a conference paper at ICLR 2021
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kemel: Convergence and gen- eralization in neural networks. In Advances in neural information processing systems, pp. 8571- 8580, 2018.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. arXiv preprint arXiv:1912.04977) 2019.
Michael Kamp and Linara Adilova. Distributed Learning Platform, 2020. https://gitie
com/ Eraunhofer--rrrr saanhoofr-iais/lppl.... form.
Michael Kamp. Linara Adilova, Joachim Sicking, Fabian Hüger, Peter Schlicht, Tim Wirtz, and Stefan Wrobel. Efficient decentralized deep learning by dynamic model averaging. In Joint European Conference Machine Learning and Knowledge Discovery in Databases, pp. 393- 409. Springer, 2018.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sehastian U Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for on-device federated learn- ing. arXiv preprint arXiv:1910.06378 2019.
Jonas Kohler, Hadi Daneshmand, Aurelien Lucchi, Thomas Hofmann, Ming Zhou, and Klaus Neymeyr. Exponential convergence rates for batch normalization: The power of length-direction decoupling in non-convex optimization. In The 22nd International Conference on Artificial Intel- ligence and Statistics, pp. 806-815, PMLR, 2019.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton. Imagenet classification with deep convo- lutional neural networks. In Advances in neural information processing systems, pp. 1097-1105, 2012.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-2324, 1998.
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine, 37(3):50-60, 2020a.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. In Conference on Machine Learning and Systems, 2020a, 2020b.
Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang. On the convergence of fedavg on non-iid data. arXiv preprint arXiv: 2019.
Yanghao Li, Naiyan Wang. Jianping Shi, Jiaying Liu, and Xiaodi Hou. Revisiting batch normaliza- tion for practical domain adaptation. arXiv preprint arXiv:1603.. 047779 2016.
Yanghao Li, Naiyan Wang, Jianping Shi, Xiaodi Hou, and Jiaying Liu. Adaptive batch normalization for practical domain adaptation. Patlern Recognition, 10:109-117, 2018.
Quande Liu, Qi Dou, Lequan Yu, and Pheng Ann Heng. Ms-net: Multi-site network for improving prostate segmentation with heterogeneous mri data. IEEE Transactions on Medical Imaging, 2020.
Ping Luo, Xinjiang Wang, Wenqi Shao, and Zhanglin Peng. Towards understanding regularization in batch normalization. arXiv preprint arXiv:/8099 1888808846666 2018.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficii learning of deep networks from decentralized data. In Artificial Intelli- gence and Statistics, pp.
2017.
Ari S Morcos, David GT Barrett, Neil C Rabinowitz, and Matthew Botvinick. On the importance of single directions for generalization. arXiv preprint 1rXiv:1803.06959) 2018
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Ng. Reading digits in natural images with unsupervised feature learning, 2011.

---

Published as a conference paper at ICLR 2021
Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. Proceedings of the IEEE International Conference on Computer Vision, pp. 1406-14155 2019.
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. arXiv preprint 00X1: 2020.
Amirhossein Reisizadeh, Farzan Farnia, Ramtin Pedarsani, and Ali Jadbabaie. Robust federated learning: The case of affine distribution shifts. arXiv preprint arXiv:2006.08907. 2020.
Nicola Rieke, Jonny Hancox, Wenqi Li, Fausto Milletari, Holger R Roth, Shadi Albarqouni, Spyri- don Bakas, Mathieu N Galtier, Bennett A Landman, Klaus Maier-Hein, et al. The future of digital health with federated learning. NPJ digital medicine, 3(1):1-7,200
Theo Ryffel, Andrew Trask, Morten Dahl, Bobby Wagner, Jason Maneuso, Daniel Rueckert, and Jonathan Pusserat-Palmbach. A generic framework for privacy preserving deep learning, arXiv preprint 1rXiv:181I.0400 2018.
Kate Saenko, Brian Kulis, Mario Fritz, and Trevor Darrell. Adapting visual category models to new domains. In European conference computer vision, pp. 213-226. Springer, 2010.
Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterizzttoon accelerate training of deep neural networks. In Advances neural information processing systems, pp. 901- 909, 2016
Shibani Santurkar, Dimitris Tsipras, Andrew Ilyas, and Aleksander Madry. How does batch nor- malization help optimization? In Advances in Neural Information Processing Systems, pp. 2483- 2493, 2018
Jan van den Brand, Binghui Peng, Zhao Song, and Omri Weinstein. Training (overparametrizeed) neural networks in near-linear time. arXiv e-prints, pp. MeXiv-2006.. 2020.
Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Federated learning with matched averaging. arXiv preprint arXiv:2002.064440 2020.
Xinwei Zhang, Mingyi Hong, Sairaj Dhople, Wotao Yin, and Yang Liu. Fedpd: A federated learning framework with optimal rates and adaptivity to non-iid data. arXiv preprint arXiv:2005. 1418, 2020.
Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data, arxiv preprint arXiv: 00582, 2018
12

---

Published as a conference paper at ICLR 2021
APPENDIX
Roadmap of Appendix The Appendix is organized as follows. We list the notations table in Section We provide theoretical proof of convergence in Section B. The algorithm of FedBN is described in Section C. The details of experimental setting are in Section D and additional results on benchmark datasets are Section E. We show experiment on synthetic data Section F. We demonstrate the ability of generalizing FedBN to test on new client in Section G.
A NOTATION TABLE
Notations Description
features, IRd dimension of x labels, E R probability distribution total number of clients total number of epochs in training number of local itcration in FL
d
y
P(-)
N
T
E
M
number of training samples each client
[N]
set of numbers, [N] [I....,N} indicator for client, i E [N]
j
indicator for sample in each client, j E [M]
(x)) w5)
the j-th training sample in client i
m
number of neurons in the first layer indicator for neuron, k E [m]
k
Ve
parameters for the k-th neuron in the first layer
v
Is
vector norm, || v Ils A covariance matrix for features in client i, S; T
given matrix S
S;
P.q
indicator for sample,
E [NM]
f
two layer ReLU neural network with BN two layer ReLU neural network with BN with client-specified BN parameters
f*
V
parameters of the first phase neurons, V Rmxd
Y
the scaling parameter of BN top layer parameters of the network ReLU activation function, of.) = max(..0) Gaussian with mean μr and covariance Σ Rademacher distribution
N
1.1
a
variance of vı at initialization empirical risk with square loss for network f evolution dynamic for FedAvg at epoch
I(f)
A(
V(
evolution dynamic with respect to V for FedAvg at epoch
G(
evolution dynamic with respect to for FedAvg at epoch t
A
(
evolution dynamic for FedBN at epoch / evolution dynamic with respect to V for FedBN at epoch evolution dynamic with respect to r for FedBN epoch t the minimal eigenvalue of matrix A expectation of G
G
Amin
(A)
G
G 8000
expectation of (f)
Table 2: Notations occurred in the paper.
13

---

Published as a conference paper at ICLR 2021
B CONVERGENCE PROOF
B.1 EVOLUTION DYNAMICS
In this section, we calculate the evolution dynamics (t) for training with function f and A for training with f*. Since the parameters are updated using gradient descent, the optimization dynamics of parameters are
dv OL
d aL dm
dt OVR dt
Let f. = f(x()) ). Then, the dynamics of the prediction of the p-th data point in site ip is
0fp
m
0fp dvk
0fp dri
m2
Ofp ØL
m
Ofp OL
Öt
OV dt Oy
OVa Ovk
8nk dym
The gradients of fp and L with respect to Va and 7h are computed as Tn
0f (t)
Cp
(1) (t)(t),
Ovk
mm
V[v (t)lls)s
aL (t) =
NM
Cx (t()) (0)
Σ (f₂(t) Ua)
[iv(t)]ls,, x4
(t),
Ovk
0fr (f) = m v[ (t)lls,
C₁
o (va(t) (xp)t))>p).
0yk
AL (t)
vm Σ (f)
yn)
Ck [vc(t)]ls,,
(vx(t)Txx xxy)
07
where fp 1(x)) Xp
S. uu Iull
jx, and pl)(1) 11 (xx()))xxx0)))
We define Gram matrix Vtt) and Gtt) as
m
(o) x2((0=
V (1)
(acr - nu (1)) (v)Is, (1)a()) xp
Xy
(na(1) (6)
m
(t) m
(t)ls ([v(0) (vx(t) xp) (va(t) x,)
It follows that
NM V pq (t)
NM Σ (t) -yy)
(t(()=
a² (fg(t) vg), TZ(t)
Let f = (fi...,, fn) (x1). , (XNM)) T. The full evolution dynamic is given by df
dt = --(((((((t) - y), where A(t) = V(t)
a + G(().
Similarly, we compute Gram matrix V1( and G1 for FedBN with f as
(0)
(1) (acx)2 Thi (t) (6) [t)()) [[)]())
xp
(z)
(8)
G (1) 1 m
Σ 2 (1) s. [[ve()) a (vx(1) xx) (v=(1)) x,) 1fip = 1g).
m
Thus, the full evolution dynamic of FedBN is
dr dt
V*
- y), where A A² =
a + (t)
14

---

Published as a conference paper at ICLR 2021
B.2 PROOF OF LEMMA 4.3
Dukler et al. (2020) proved that the matrix G is strictly positive definite. In their proof, G is the covariance matrix of the functionals Op define
0p(v) = xp)
over the Hilbert space V of L2 (N (0, II). oe is strictly positive definite by showing that Φ1, ... ONN are linearly independent, which is equivalent to that
01 + c202 + +... CNMÓNM in
(10)
holds only for c 0 for all p.
Let G. denote the i-th M x M block matrices on the diagonal of G Then we have
G diog( (G , G
To prove that t is strictly positive definite, will show that G is positive definite. Let us define
(v) := (v"xx) x,) 1{j site i}, j= 1...... M.
Then, going show that
c + 0202, + ... CM CM
(11)
holds only for oj 0, E [M]. Suppose there exist c cas that not identically 0, satisfying (11). Let the coefficients for client be c1, CM and let the coefficients for other client be 0. Then, have sequence of coefficients satisfying (10), which is contradiction with that G strictly positive definite. This implies G is strictly positive definite. Namely, G eigenvalues are positive. Since the eigenvalues of *00 are exactly the union of the eigenvalues of G Amin *00 is positive and thus, G is strictly positive definite.
B.3 PROOF OF COROLLARY 4.6
To compare the convergence rates of FedAvg and FedBN when E = 1, we compare the exponential
factor in the convergence rates, which are (1 nua/2) and (1 m/ /2) for FedAvg and FedBN, respectively. Then, it reduces to comparing μu = Amin G and μr Amin(G" (G Comparing
equation (7) and (9), G takes the M x M block matrices the diagonal of oo
G
G
GIN
G 0
G
G
G
G
G
G
G
where G the i-th x M block matrices on the diagonal of G By linear algebra,
Amin (G 2 Amin (G Vi [N].
Since the eigenvalues of "00 exactly the union of eigenvalues of G have
Amin a0)) min Amin (G,)))
<E(N]
2 Amin (
Thus, (1 (1-1///2) 2 -1-11 (1/5//2) and we can conclude that the convergence rate of FedBN is faster than the convergence of FedAvg.
15

---

Published as a conference paper at ICLR 2021
C FEDBN ALGORITHM
We describe the details algorithm of our proposed FedBN as following Algorithm 1:
Algorithm Federated Learning using FedBN
Notations: The user indexed by k, neural network layer indexed by I, initialized model param- (1)
eters: "okk local update pace: E, and total optimization round T.
1: for each round = 1, .,, T do 2: for each user and each layer / do
(1 1111 t (GD(w
3:
4: end for
5: if mod(t, E) = 0 then
6:
for each user k and each layer do
7:
if layer 7 is not BatchNorm then
(1 1+1 + t Σx w t+1
8:
9: end
10: end for
11: end
12:
end for
16

---

Published as a conference paper at ICLR 2021
D
EXPERIMENTAL DETAILS
D.1
VISUALIZATION OF BENCHMARK DATASETS
2
MNITT
SVHN
USPS
density pix)
20-
Syrthhiggts
USPS
SANN
MNIST MM
1
Synthlgits
MAIST .M
0.0
1.0
pixel value
(2)
(b)
Figure 6: Data visualization. (a) Examples from each dataset (client). (b) Non-iid feature distribu- tions across the datasets (over random 100 samples for each dataset).
We show image examples from the five benchmark datasets and the pixel value histogram. It ob- viously presents the heterogeneous appearances and shifted distributions. Clients formed from five benchmark datasets are viewed as non-iid.
D.2 MODEL ARCHITECTURE AND TRAINING DETAILS ON BENCHMARK
We illustrate our model architecture and training details of the digits classification experiments in this section.
For our benchmark experiment, we use six-layer Convolutional Neural Network (CNN) and its details are listed Table
Model Architecture
Layer
Details
Conv2D(3, 64, 5, 1, 2)
BN(64), ReLU, MaxPool2DC2 2)
Conv2D(64, 64, 1, 2)
BN(64), ReLU, MaxPool2D(2,, 2)
Conv2D(64, 128, 5, 1, 2) BN(128), ReLU Conv2D(64, 128, 5, I, 2)
BN(128), ReLU FC(6272, 2048)
BN(2048), ReLU
FC(2048, 512)
BN(512), ReLU
FC(512, 10)
Table 3: Model architecture of the benchmark experiment. For convolutional layer (Conv2D), list parameters with sequence of input and output dimension, kernal size, stride and padding. For max pooling layer (MaxPool2D), list kernal and stride. For fully connected layer (FC), list input and output dimension. For Barmaliormaiization layer (BN), we list the channel dimension.
17

---

Published as a conference paper at ICLR 2021
Training Details. We give detailed settings for the experiments conducted in 5.1: (1) convergence rate (Table 4), (2) analysis local update epochs (Table 5), (3) analysis of local dataset size (Table 6), (4) effects of statistical heterogeneity (Table 7) and (5) comparison with ofte-of-theeart (Table
8). Each table describes the number of clients, samples and the local update epochs.
SGD optimizer with learning we set batch size to 32 and training epochs to 300. For hyper-parameter μ4, we use the best value y 10-2 founded by grid search from the the default settings in FedProx Li al. (2020b).
During training process,
10-2 and cross-entropy loss,
Datasets
SVHN
USPS
SynthDigits
MNIST-M
MNIST
Number of clients
1
Number of samples
743
743
743
743
743
Local update epochs
1
Table 4: Settings for convergence rate. Each dataset has client with 743 samples, local update epoch is set 1.
Datasets
SVHN
USPS
SynthDigits
MNIST-M
MNIST
Number of clients
1
Number of samples
743
743
743
743
743
Local update epochs
1,4.8,16
1,4,8,16
1,4,8,16
1,4,8.16
1,4,8,16
Settings for local update epochs. Each dataset has client with 743 samples, local update epoch for all datasets set to 4, 8, 16 successively.
Table 5:
18

---

Published as a conference paper at ICLR 2021
Datasets
SVHN
USPS
SynthDigits
MNIST-M
MNIST
Number of clients
1
Number of samples
Local update epochs
Table 6: Settings for local dataset size, set local update epochs to and each dataset has client.
The number of samples (74,371, 743, 1487, 2975, 4462, 7438).
Datasets
SVHN
USPS
SynthDigits
MNIST-M
MNIST
Number of clients
[1. 10]
10]
[1. 10]
10]
[1. 10]
Number of samples
[1, 10]x743
[1, 143
[1. 10]
[1. 103
[1. 10]x743
Local update epochs
Table 7: Settings for statistical heterogeneity, [1, 10] for the range from to 10. We increase number of clients step by step and number of samples will increase accordingly.
Datasets
SVHN
USPS
SynthDigits
MNIST-M
MNIST
Number of clients Number of samples
1
1
743
743
743
743
743
Local update epochs
1
Table 8: Settings for comparison with SOTA, we use epoch for comparison experiment.
client with 743 samples and local update
19

---

Published as a conference paper at ICLR 2021
D.3 MODEL ARCHITECTURE AND TRANING DETAILS IMAGE CLASSIFICATION TASK ON OFFICE- CALEEECCALECHHOOO AND DOMAINNET
In this section, provide the details of our model and training process on both 1ffice-Caltech10 Gong et al. (2012) and DomainNet Peng et al. (2019) dataset.
Model Architecture For the image classification tasks on these two real-worldd datasets Office- Ca and DomainNet data, we use adapted AlexNet added with BN layer after each convolu- tional layer and fully-connected layer (except the last layer), architecture is shown in Table
Layer
Details
Conv2D(3, 64, 11, 2)
BN(64), ReLU, MaxPool2D(332 2)
C6n4 192, 5, 1, 2) BN((92), ReLU, Max 2)
2
Conv2D(64, 128, 5, 1, BN(128), ReLU Conv2D(I92, 384, 3, 1)
3
BN(384), ReLU Conv2D(384 256, 3, 1)
BN(256), ReLU
Conv2D(256, 256, 3, 1)
5
BN(256), ReLU, MaxPoll2D(3.2) 2)
6
AdapiveevvgPool2266..0 6)
FC(9216, 4096)
BN(4096), ReLU
FC(4096, 4096)
BN(4096), ReLU
FC(4096, 10)
Table 9: Model architecture for Office-Caltech10 and DomainNet experiment. For convolutional layer (Conv2D), we list parameters with sequence of input and output dimension, kernal size, stride and padding. For max pooling layer (MaxPool2D), we list kernal and stride. For fully connected layer (FC), we list input and output dimension. For BatchNormalization layer (BN), we list the channel dimension.
Training Details. Office-Caltede selects 10 common objects in Office-31 Saenko et al. (2010) and Caltech-256 datasets Griffin et al. (2007). There are four different data sources, one from Caltech-256 and three from Office-31, namely Amazon(images collected from online shopping web- site), DSLR and Webcam(images captured in office environment using Digital SL. camera and web camera).
We first reshape input images the two dataset into 256x256x3, then for training process, we use cross-entropy loss and SGD optimizer with learning rate of 10 batch size is set to 32 and training epochs is 300. When comparing with FedProx, we set μ to 10 which is tuned from the default settings. The data sample number are kept into the same size according to the smallest dataset, i.e. Office-Caltech 10 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.
20

---

Published as a conference paper at ICLR 2021
D.4 ABIDE DATASET AND TRAINING DETAILS
Here we describe the real-world medical datasets, the preprocessing and training details.
Dataset:
The study was carried using resting-state IMRI (rs-IMRI) data from the Autism Brain Imaging Data Exchange dataset (ABIDE I preprocessed, (Di Martino et al., 2014)). ABIDE consortium that provides preciously collected rs-IMRI ASD and matched controls data for the purpose of data sharing in the scientific community. We downloaded Regions of Interests (ROIs) fMRI series of the top four largest sites (UM, NYU, USM, UCLA viewed as clients) from preprocessed ABIDE dataset with Configurable Pipeline for the Analysis of Connectomes (CPAC) and parcellated by Harvard- -xfordd (HO) atlas. Skipping subjects lacking filename, resulting in 88, 167, 52, 63 subjects for UM, NYU, USM, UCLA separately. Due to lack of sufficient data, 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 the dataset size in a client.
NYU
UM1
USM
UCLA1
Total Subject
167
88
52
63
ASD Subject
73
43
33
37
HC Subject
94
45
19
26
ASD Percentage
44%
49%
63%
59%
fMRI Frames
176
296
236
116
Overlapping Trunc
145
265
205
85
Table 10: Data summary of the dataset used in study.
Training Process For all the strategies, we set batch size as 100. The total training local epoch 50 with learning rate 10 with SGD optimizer. Local update epoch for each client is E 1. We selected the best parameters 0.2 in FedProx through grid search.
21

---

Published as a conference paper at ICLR 2021
E MORE EXPERIMENTAL RESULTS ON BENCHMARK DATASETS
E.1 CONVERGENCE COMPARISON OVER FEDAVG AND FEDBN
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 our theoretical analysis in Section and experimental results in Section 5.
1.4
E=1
0.6
E=4 4
E=8
1.2
0.4
E=16
1.0
0.2
Training loss
0.8
0.0
20
40
60
0.6
0.4
FedAvg
0.2
FedBN
0.0
25
50
75
100 Epochs
125
150
175
200
Figure 7: Training loss epochs with different local update frequency.
E.2 DETAILED STATISTICS OF FIGURE 5
In Figure 5, we compare the performance with respect to accuracy of FedBN and alternative meth- ods. We show the detailed accuracy in the following Table 11.
Methods
SVHN
USPS
Synth
MNIST-M
MNIST
Single
65.25 (1.07)
95.16 (0.12)
80.31 (0.38)
77.77 (0.47)
94.38 (0.07)
FedAvg
62.86 (1,49)
95.56 (0.27)
82.27 (0.44)
76.85 (0.54)
95.87 (0.20)
FedProx
63,08 (1.62)
95.58 (0.31)
82.34 (0.37)
76.64 (0.55)
95.75 (0.21)
FedBN
71.04 (0.31)
96.97 (0.32)
83.19 (0.42)
78.33 (0.66)
96.57 (0.13)
Table 11: The detailed statistics reported with format mean (std) of accuracy presented Fig. 5
22

---

Published as a conference paper at ICLR 2021
E.3 COMPARE FEDBN WITH CENTRALIZED TRAINING
To better understand the significance of the numbers reported in our main context, we compare FedBN with centralized training, that pools all training data in center. We present the testing accuracy on each digit dataset in Table 12. FedBN, federated learning with data-specific BN layers, could achieve comparable performance with vanilla centralized training strategy.
SVHN
USPS
Sigtt
MNIST-M 79.65 (0.24)
MNIST 96.53 (0.19)
Centralized
74.18 (0.44)
96.46 (0.30)
84.57 (0.38)
FedBN
71.04 (0.31)
96.97 (0.32)
83.19 (0.42)
78.33 (0.66)
96.57 (0.13)
Table 12: Testing accuracy each testing sets with format mean(std) from 5-trial run.
E.4 DIFFERENT COMBINATIONS OF AND
In this section, show different combinations of local update epochs E and batch size B. Specif- ically, E E {1. 4, 16} and E {10, 50, )o). oo denotes full batch learning. Following the setting in original FedAvg paper McMahan et al. (2017),
present the comparisons between FedBN and FedAvg on each combination of and B in Table 13. The results are good agreement that FedBN consistently outperform FedAvg and robust to batch size selection. Further, depicts the test sets accuracy vs. local epochs under different combination of E' and B in Figure 8.
Setting
SVHN
USPS
SynthDigits
MNIST-M
MNIST
FedAvg
65.50
97.04
84.25
81,65
96.55
B=10, E=1
FedBN
76.18
97.37
86.51
82.81
97.41
FedAvg
69.80
96.67
85.63
82.54
97.21
B=10, E=4
FedBN
76.23
97.04
86.99
83.14
97.05
FedAvg
65.05
95.05
83.74
80.79
96.71
B=10, E=16
FedBN
75.56
96.13
84.78
80.29
96.44
FedAvg
62.42
95.32
81.66
75.28
96.06
B=50, E=1
FedBN
70.70
97.04
82.74
78.38
96.57
FedAvg
61.67
95.16
80.69
74.44
95.71
B=50, E=4
FedBN
69.85
97.10
81.78
77.56
96.40
FedAvg
60.00
94.68
79.37
73.39
95.28
B=50, F=16
FedBN
67.67
96.94
80.39
76.54
95.66
FedAvg
60.99
94.57
79.69
74.36
95.86
B=00, E=1
FedBN
65.98
96.29
79.75
76.79
96.15
FedAvg 59.07
95.38
79.88
73.97
94.51
B=o0, E=4
FedBN
65.25
96.34
79.99
73.96
95.51
FedAvg 61.88
94.68
78.69
74.36
95.46
B=co, E=16
FedBN
64.39
95.16
78.22
73.96
95.57
Table 13: Test accuracy using different combinations of batch size B and local update epoch E on benchmark experiment with the default non-iid setting.
23

---

Published as a conference paper at ICLR 2021
0.9
0.8
Accuracy (%6
0.7
030
046
036
032
0.6
E-1 8-10
640
6.8
E=1 8
807
016
0.5
E-4 8-30
E-16 B-10
E-16 B-50
E-16 B-
0.4
100
200
300
400
500
Epochs
Figure 8: Test set accuracy curve (average of 5 datasets) of using different local updating epochs E and batch size B for FedBN.
E.5 DETAILED STATISTICS OF VARYING LOCAL DATASET SIZE EXPERIMENT
Considering putting all results in figure (15 lines) might affect readability of the figure, excluded the statistics of FedAvg in our Fig. 4 (b), the ablation study of our method on the effect of local dataset size. Here, list the full results in Table 14. It is not too surprising that Singleset can be the best when the local client gets a lot of data.
Setting
100%
60%
40%
20%
10%
5%
1 1%
SingleSet
98.09
97.84
97.22
96.14
94.35
90.86
75.28
MNIST
FedAvg
98.96
98.54
98.13
97.51
96.22
93.79
79.94
FedBN
98.91
98.63
98.34
97.78
96.72
94.67
85.22
SingleSet
85.74
84.62
82.75
76.42
66.81
52.62
12.06
SVHN
FedAvg
82.08
79.56
77.37
70.84
63.80
49.15
23.67
FedBN
86.93
84.72
82.87
78.20
71.31
61.53
31.98
SingleSet
98.87
98.49
97.85
96.94
95.11
93.01
80.11
USPS
FedAvg
98.33
97.85
97.42
96.61
95.59
93.76
79.09
FedBN
98,82
98.92
98.55
98.17
97.58
96.24
85.05
SingleSet
94.33
92.82
91.02
86.77
80.47
70.61
14.10
Synth
FedAvg
93.98
92.57
91.04
87.03
82.17
72.76
42.11
FedBN
94.40
92.81
91.75
88.03
83.06
74.85
43.76
SingleSet
93.30
91.63
89.41
84.34
77.59
66.02
17.23
MNISTM
FedAvg
90.59
88.91
86.21
82.11
76.93
67.97
41.71
FedBN
91.35
89.95
87.79
83.73
78.80
70.04
44.17
Table 14: Model performance varying dataset sizes local clients
E.6 TRAINING UNEQUAL DATASET SIZE
In our benchmark experiment (Section 5.1), truncate the sample size of the five datasets to their smallest number. This data preprocessing intends to strictly control non-related factors (e.g., imbal- anced sample numbers across clients), that the experimental findings more clearly reflect the effect of local BN. In this regard, truncating datasets is a reasonable way to make each client have equal number of data points and local update steps. It is also possible to keep the data sets in their
24

---

Published as a conference paper at ICLR 2021
original size (which is unequal), by allowing clients with less data repeat sampling. In this way, all clients use the same batch size and same local iterations of each epoch. We add results of such setting with 10% and full original datasize in Table 15 and Table 16 respectively. It is observed that FedBN still consistently outperforms other methods.
SVHN
USPS
SynthDigits 39116
MNIST-M
MNIST
Method 7943
743
5600
5600
FedAvg 87.00
98.01
97.55
88.69
98.75
FedProx
86.75
97.90
97.53
88.86
98.86
FedBN 89.34
98.28
97.83
90.34
98.89
Table 15: Testing accuracy of each clients when clients' training samples are unequal using 10% of original data. The number of training samples for each client denoted under their names.
SVHN
USPS
SynthDigits
MNIST-M
MNIST
Method
79430
7430
391 160
56000
56000
FedAvg 99.59
92.27
98.71
99.30
95.27
FedProx
99.50
92.12
98.66
99.27
95.44
FedBN 99.62
94.34
98.92
99.54
96.72
Table 16: Testing accuracy of each clients when clients' training samples are unequal using full size data. The number of training samples for each client are denoted under their names.
25

---

Published as a conference paper at ICLR 2021
F SYNTHETIC DATA EXPERIMENT
Settings We generate data from two-pair of multi-Gaussian distributions. For one pair, samples (x and (ær, 1) are sampled from (--1, Σn
and (11, 21) respectively, with coveriance 2n R10x10 For another pair, samples (z. 0) and (3. 1) are sampled from W(-1, Σ2) and (11, Σg) respectively, with coveriance Σg 10 10 Specifically, we design convariance matrix Ej as an identity diagonal matrix and Σg is different from Σ1 by having non-zero values on off-diagonal entries. We train two-layer neural network with 100 hidden neurons for 600 steps using cross- entropy loss and SGD optimizer with 1 x 10 learning rate. Denote W; and by are the in-connection weigths and bias term of neuron We initialize the model parameters with We NO, II), by (0, a²), where
10.
Results. The aim of synthetic experiments to study the behavior of using FedBN with a controlled setup. We achieve 100% accuracy on binary classification for FedAvg and FedBN. Fig. 9 shows comparison of training loss curve over steps using FedAvg and FedBN, presenting that FedBN obtains significantly faster convergence than FedAvg.
0.6
FedAvg FedBN
FedAvg FedBN
0.5
0.4
Training
0.3
Training
0.3
0.2
0.2
0.1
0.1
0.0
0.0
100
200
300
400
500
600
100
200
300
400
500
600
Steps
Steps
(a) client
1
(b) client 2
Figure 9: Training loss on synthetic data. Data client is generated from Diagonal Gaussian, client is generated from combination Diagonal Gaussian and Full Gaussian.
26

---

Published as a conference paper at ICLR 2021
G TRANSFER LEARNING AND TESTING ON UNKNOWN DOMAIN CLIENT
In this section, we discuss out-of-domain generalization of FedBN and prove the solutions for the following two scenarios: I) transferring FedBN new unknown domain clients during training: 2) testing an unknown domain client.
If 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.
Testing the global model on a new client with unknown statistics outside federation requires allow- ing access local BN parameters at testing time (though BN layers are not aggregated 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 own data. Such 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. Specifically, Morpho-glow -looo-llolllaa and Morpho-local from Morpho-MNISS (Castro et al., 2019) as the two new clients. The new clients contain substantially perturbed digits. Specifically, Morpho-global containing thin- ning 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.
use the digits classification task and treat the two unseen datasets
Morpho-global
Morpho-local
FedBN
92.45
94.61
FedProx FedAvg 91.28
92.35
94.31
93.55
Table 17: Generalizing the global model to unseen-domain clients.
27


Reviews:
[{"review_id": "ugDU0fozHAW", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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). \n\nThis paper is self-contained in the sense that the initial problem is well-defined with some theoretical background and visualisations. Then a very simple solution is validated both theoretically and empirically. While the underlining idea is very, very simple, the theoretical and empirical justifications remained unclear until this work. According to this paper, we now are able to motivate the use of local BN in the context of FL. Moreover, this approach has been evaluated (convergence rate, variations w.r.t local epochs, different data distributions, heterogeneity variations, results vs SOTA)  in different non-IID FL tasks (USPS, MNIST, SynthDigits, MNIST-H). Finally, an example on \"real-world\" data is given to better highlight the importance of the demonstration. The proposed FedBN convincingly outperformed all the other approaches in all the setups. \n\nHowever, I found the paper hard to read due to the lack of a clear plan. I think the content should be a bit more structured. \n\nSome remarks and concerns are listed bellow:\n\nPros:\n+ Good theoretical backup for a very simple idea.\n+ Finally validating the use of local BN that was certainly used by many researchers but for \"unexplained\" reasons.\n+ Good results compared to SOTA methods.\n+ Real-world datasets are not really common in FL papers. It's great to see one. \n\nCons:\n- The idea remains a bit \"too\" simple, can't we build on the top of these findings to further improve the BN ?\n- Where are the appendix ?\n- I find the organisation of the Sections a bit hard to follow. In particular, from the end of page 3 to the end of page 5 the reader has to follow a single big Section without any segmentation while others sections are 2 paragraphs long. As an example: Is Section 3 really needed ? Can't it be integrated into a revised Section 4 (with proper sub-sections) ?\n- As BatchNorm statistics are kept local to each client. What happen at testing time ? Let's say we want to distribute the trained model to a new client that does not have already computed statistics ?\n\nRemarks:\n\nThe paper could benefit from some proof-reading. Examples are:\n- Colloray 4.6\n- In particular, the average model w = (w⇤ +w⇤)/2 and average BN parameters \u0000 = (\u0000⇤ +\u0000⇤)/2 has a high error -> unclear\n- Widely known ag- gregation strategy in FL, FedAvg (McMahan et al., 2017) it often su↵ers when data is heterogeneous over local client. -> not correct\n- As Fig. 1 shows the local squared loss between is very di↵erent between the two clients -> not correct\n- What is ck in Eq 1 ?\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "ZTJezEskkiJ", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a minor modification to the existing FL learning framework, which can be easily implemented. The major contribution of this work comes from a few theoretical analyses. The experiments are conducted on several publically accessible data sets with intuitive explanations. \n\nPros:\n- The FL toy example is quite interesting, and we get a nice starting idea for your overall paper.\n- The theoretical analysis seems pretty solid and intuitive. I didn't fully proofread all the details, but the parts that I read look correct to me.\n- The experiments reveal certain properties of FedAve and FedBN, most results are convincing and the relative improvement looks pretty predictable based on the explanation.\n\nCons:\n- The motivation on the toy example may not work well on real-world data set, especially since these assumptions may not hold in general. \n- The experiments are oversimplified and lacks a proper explanation for this kind of preprocessing.\n- Without any large-scale publically-accessible data set being utilized, I concern a lot about its performance on real applications.\n\nDetailed concerns:\n- The necessity of preprocessing all data set into the same setting is not provided. Why only keep 7438 examples for each data set? Does it have something to do with the non-iid setting? Is it possible to keep the data set in their original size? \n- I checked the supplementary as well and didn't find how do you make sure the data set assigned to each client is non-iid.  The claimed covariate shift and concept shift didn't have corresponding experimental settings being explicitly stated, which looks quite confusing to me. Also, please add the comparison with the non-iid scenario with the one proposed in FedAvg.\n-  I highly suggest the authors add a few large-scale data set to verify the effectiveness of FedBN, e.g., cifar-10, cifar-100, and imagenet. Because these data sets are more close to the scenario where each client has non-iid data. \n- Each client has 10% data seems a bit too easy especially for some simple data set, e.g., MNIST, and its variants, I highly believe a comprehensive study like the one did in the FedAvg paper by changing K, B, E, and all their combinations, is the correct way to evaluate your proposed framework. Especially, FedBN is extremely easy to implement. Without these comparisons, I believe the experiment is not complete.\n- Code is not attached, many details of FedBN remains unclear. The results may perform differently when we implement our own version. \n\n\nMinor comments: \nInconsistent notations: sometimes FedAvg but sometimes using FedAVG.", "extra_scores": {}, "extra_text": {}}, {"review_id": "4AaSIhpxO5M", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Update following answer:\n\nThanks for your detailed answer, which confirms me in my initial assessment.\n\n------\n\n\n1/ Summary of the paper\n\nThis paper introduces the use of local batch normalisation layers in order to circumvent data shifts issues in FL, called FedBN.\nBuilding on a simplified model of BN (neural network with 1 hidden layer and BN rescaling) that was previously introduced by [1] to study the impact of BN on convergence, it is proven that the convergence of the proposed FedBN on training data can only be faster than the convergence of FedAvg on the same data points.\nParametric experiments on a simple heterogeneous dataset (built as the union of digit classification tasks) show that the proposed method yields a better performance, and a more stable one, than standard FedAvg.\nExperiments on 3 real-world datasets containing heterogeneity to simulate different centers show again that the proposed FedBN method improves upon FedAvg.\n\n2/ Acceptance decision\n\nWhile ideas related to BN and FL have been floating in the community this year, this paper proposes a novel approach with extensive and convincing numerical results, and interesting theoretical results on the effect of fedBN.\nProvided the paper is updated to better acknowledge other very related works, I think it should be accepted as it will be a very valuable contribution to the FL community.\n\n\n3/ Supporting arguments\n\nA/ Novelty\n\nA very close reference not acknowledged by the paper is [2], which builds on the same ideas of domain adaptation as this paper and proposes to keep local BN weights in the FL setting, called SiloedBN, showing that results are improved and more stable with respect to FedAvg.\nHowever, the proposed FedBN is different from siloedBN [2] in the sense that SiloedBN proposes only to keep local BN statistics, while FedBN keeps local BN layers altogether (trainable and untrainable parameters thereof).\nTherefore, although not entirely new, FedBN is still novel.\n\nB/ Theoretical results\n\nSetting aside the drift issues encountered when the number of local updates increases, the paper uses the same formalism as [1] to get a linear convergence bound for FedAvg on the training set, which is controlled by some key constant mu_0 (the higher, the faster the convergence).\nFor this analysis, the contribution of this paper is to put FedBN under the same framework as [1] and to show that the resulting constant mu_0^* is larger than the constant mu_0: therefore, FedBN’s training error diminishes faster than FedAvg’s.\nAlthough the underlying model is a simplification of the experimental reality, the problem tackled is very complex, so this result is still a significant contribution for the community.\n\nC/ Experimental results\n\nThe experimental results are in three parts: \n- An introductory toy example, which helps to better understand the importance of local BN layers in the case of domain shift.\n- Benchmark on a simple hand-crafted FL digit recognition dataset built by concatenating different digit datasets with different domain shifts (e.g. MNIST, SVHN…), thereby borrowing from the domain adaptation literature. While I have some remarks to make the results even more complete (cf next section), I think these results help to better understand the behaviour of FedBN and FedAvg under heterogeneity.\n- Experimental results on 3 hand-crafted FL datasets with a varying number of centers (4, 6) and different tasks. In almost all cases, FedBN significantly improves the final testing accuracy upon FedAvg and FedProx in almost centers.\nThese experiments are exhaustive and convincing, and are definitely an asset of this submission.\n\nD/ Writing\nAlthough dense, the paper is well written and easy to follow. There are minor typos.\n\n4/ Additional comments\n1. How could FedBN be adapted to transfer a model on a new center from another domain?\n2. There is a minor error in the proof of Corollary 4.6, which does not invalidate the results. Indeed, the authors claim that since G^{*, \\infty} = diag(G_1^{\\infty}, \\ldots, G_N^{\\infty}), one has, for all i, \\lambda_min(G^{*, \\infty}) > \\lambda_min(G_i^{\\infty}). This is false in general: one can only claim that \\lambda_min(G^{*, \\infty}) \\geq \\min_i \\lambda_min(G_i^{\\infty}). Since all the minimal eigenvalues of G_i^{\\infty} are lower bounded by the minimal eigenvalues of G^{\\infty}, the final result still holds.\n3. This is extremely minor, but in Corollary 4.6, I am not sure that in all generality one can have strict inequalities: only >= statements may be available\n4. Related to the previous remark, although very interesting, the result of Corollary 4.6 is frustrating as it does not quantify the speed improvement. Even if a closed-form quanitification is out of hand, it may be interesting to produce numerical experiments on synthetic data on this simplified model to quantify the gap between the convergence bounds.\n5. In section 5.1, how is computed the testing accuracy reported in figure 4? Is it the mean of the per-center accuracy on the local testing datasets? The training sets are artificially balanced, but how balanced are the testing datasets?\n6. In order to better understand the significance of all the numbers reported, it would have been nice to report the performance of a model trained in a pooled-equivalent fashion, i.e. with all data points in a single center.\n7. In Fig 4 b), one studies the effect of the local dataset size on the performance, and, not surprisingly, the performance of all methods diminish in this case. Why isn’t FedAvg reported here? In particular, why was the 10% fraction chosen in the other experiments, even if it led to suboptimal results?\n8. In Figure 4 c), are the points with different x values related in any case? Or are they coming from independent experiments. The text « we started with including… then, we simultaneously added n clients » is a bit ambiguous.\n\nTypos and other remarks\n9. The use of both m and M in Sec 4 can be confusing\n10. In Lemma 4.3, assumption 1 should be replaced by assumption 4.1\n11. There are 2 typos in theorem 4.4: missing parentheses around Equation 2, and n should be replaced by N.\n12. In Figure 5, the colours are hard to parse for colour-blind people.\n\nRefs\n——\n[1] Dulker, Gu, and Mont\\’{u}far, Optimization theory for relu neural networks trained with normalisation layers, in proceedings of ICML 2020\n[2] Andreux, Ogier du Terrail, Beguier, and Tramel, Siloed federated Learning for Multi-Centric Histopathology Datasets, in proceedings of MICCAI DCL 2020", "extra_scores": {}, "extra_text": {}}, {"review_id": "1iOoO327bM2", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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. \n\nMerits:\n1. The paper tackles an important problem in federated learning, which may have practical impact.\n2. There is an effort on quantifying the performance gain of FedBN vs FedAvg.\n\nCons:\n1. The assumption that the feature shift on local datasets is captured by the difference in local covariances needs better justification. Why this is a valid assumption?\n\n2. The comparison is on the convergence to f and f* under FedAvg and FedBN. It is not clear if f and f* have the same learning accuracy or training loss. If not, then the comparison is not fair.\n\n3. The convergence result is not stated in a clear form. In general, exact linear convergence to the optimal solution is not possible for SGD-type algorithm. The condition under which this holds true needs to be explicitly mentioned in Theorem 4.4 and Corollary 4.5. In addition, how the convergence depends on the averaging interval E is not clear. In FedAvg, there is non-negligible error in non-i.i.d. setting that depends on E. What technique you have used to eliminate it? These details are all unclear. \n\n4. The current analysis is mainly following Dukler et al. (2020). The new technical challenge and difficulty in analyzing the federated version of Dukler et al. (2020) needs to be highlighted. \n\n5. The simulations are not convincing. On one side, the state-of-the-art baseline FedAdam Reddi et al. (2020) has not been compared; on the other hand, the performance gain relative to FedAvg seems to be marginal. ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper addresses the problem of batch normalization (BN) in federated learning, which is of great interest to the community including practitioners.  The proposed method here simply excludes the BN parameters from the aggregation, and evolves them locally.

As a main contribution, reviewers particularly liked the solid justification of the proposed scheme, both with substantial theory and extensive experiments. Presentation style can be slightly improved, the usage at test time can be clarified more, and some references mentioned by R3 should be added, but this overall does not affect the strong level of contributions present in the work, and the discussion phase with the authors was already constructive.
Decision: Accept (Poster)