FedBE: Making Bayesian Model Ensemble Applicable to Federated Learning
Authors: Hong-You Chen
Wei-Lun Chao
Venue: ICLR 2021 Poster (2021)

Abstract: Federated learning aims to collaboratively train a strong global model by accessing users' locally trained models but not their own data. A crucial step is therefore to aggregate local models into a global model, which has been shown challenging when users have non-i.i.d. data. In this paper, we propose a novel aggregation algorithm named FedBE, which takes a Bayesian inference perspective by sampling higher-quality global models and combining them via Bayesian model Ensemble, leading to much robust aggregation. We show that an effective model distribution can be constructed by simply fitting a Gaussian or Dirichlet distribution to the local models. Our empirical studies validate FedBE's superior performance, especially when users' data are not i.i.d. and when the neural networks go deeper. Moreover, FedBE is compatible with recent efforts in regularizing users' model training, making it an easily applicable module: you only need to replace the aggregation method but leave other parts of your federated learning algorithm intact.

Published as a conference paper at ICLR 2021
FEDBE: MAKING BAYESIAN MODEL ENSEMBLE APPLICABLE TO FEDERATED LEARNING
Hong- You Chen The Ohio State University, USA chen. 9301 @ 8osu.ue edu
Wei-Lun Chao
The Ohio State University, USA chao 2096 oow.ousueeo. edu
ABSTRACT
Federated learning aims to collaboratively train a strong global model by accessing users' locally trained models but not their data. A crucial step is therefore aggregate local models into a global model, which has been shown challenging when users have non-i.idd data. In this paper, we propose a novel aggregation algo- rithm named FEDBE, which takes a Bayesian inference perspective by sampling higher-quality global models and combining them via Bayesian model Ensemble, leading much robust aggregation. We show that an effective model distribution can be constructed by simply fitting Gaussian or Dirichlet distribution to the local models. Our empirical studies validate FEDBE's superior performance, especially when users' data are not i.i.d. and when the neural networks go deeper. Morcover, FEDBE is compatible with recent efforts in regularizing users' model training, making it an easily applicable module: you only need to replace the aggregation method but leave other parts of your federated learning algorithm intact.
INTRODUCTION
Modern machine learning algorithms are data and computation hungry. therefore desired to collect as many data and computational resources as possible, for example, from individual users (e.g., users' smartphones and pictures taken on them), without raising concerns in data security and privacy. Federated learning has thus emerged promising learning paradigm, which leverages individuals' computational powers and data securely by only sharing their locally trained models with the server - jointly optimize global model (Konččy et al., 2016 Yang et al. 2019).
Federated learning (FL) generally involves multiple rounds of communication between the server and clients (i.e., individual sites). Within each round, the clients first train their own models using their own data, usually with limited sizes. The server then aggregates these models into a single, global model. The clients then begin the next round of training, using the global model as the initialization.
We focus on model aggregation, one of the most critical steps in FL. The standard method is FEDAVG (McMahan al. 2017), which performs element-wise average over clients' model weights. Assuming that each client's 's data are sampled i.i.d. from their aggregated data, FEDAVG has been shown convergent to the ideal model trained in centralized way using the aggregated data (Zinkevichl ct al. 2010: McMahan et al. 2017; Zhou & Cong 2017). Its performance, however, can degrade drastically such an assumption does not hold in practice (Karimireddy et al.] 2020 Li al. 2020b Zhao et al. 2018): FEDAvG simply drifts away from the ideal model. Moreover, by only taking weight average, FEDAvG does not fully utilize the information among clients (e.g.. variances), and may have negative effects on weerparamettrized models like neural networks due to their
germutation-invarian property in the weight space (Wang et al. 2020 Yurochkin et al. 2019).
To address these issues, we propose novel aggregation approach using Bayesian inference, inspired by (Maddox et al., (2019). Treating each client's model as a possible global model, we construct a distribution of global models, from which weight average (i.e., FEDAVG) is one particular sample and many other global models can be sampled. This distribution enables Bayesian model ensemble robust prediction. We
aggregating the outputs of wide spectrum of global models for show that Bayesian model ensemble can make more accurate predictions than weight average at a single round of communication, especially under the non i.i.d. client condition. Nevertheless, lacking

---

Published as a conference paper at ICLR 2021
single global model that represents Bayesian model ensemble and can be sent back to clients, Bayesian model ensemble cannot directly benefit federated learning in a multi-round setting.
We therefore present FEDBE, a learning algorithm that effectively incorporates Bayesian model Ensemble into federated learning. Following (Guha et al. 2019), we assume that the server has access to set of unlabeled data, on which we can make predictions by model ensemble. This assumption can easily be satisfied: the server usually collects its own data for model validation, and collecting unlabeled data is simpler than labeled ones. (See /See|section 6 for more discussion, including the privacy concern.) Treating the ensemble predictions as the "pseudo-labels" of the unlabeled data, we can then summarize model ensemble into a single global model by knowledge distillation (Hinton et al., 2015) using the predicted labels (or probabilities or logits) as the teacher to train a student global model. The student global model can then be sent back to the clients to begin their next round of training
We identify one key detail of knowledge distillation in FEDBE. In contrast to its common practice where the teacher is highly accurate and labeled data are accessible, the ensemble predictions in federated learning can be relatively nois?" P To prevent the student from over-fitting the noise, we apply stochastic weight average (SWA) (Izmailov et al. 2018) in distillation. SWA runs stochastic gradient descent (SGD) with cyclical learning rate and averages the weights of the traversed models, allowing the traversed models to jump out of noisy local minimums, leading to a more robust student.
We validate FEDBE CIFAR- (Krizhevsky al. 2009) and Iing-ImaeeNet (Le & Yang 2013) under different client conditions (i.e., 1.i.d. and non-i.i.d ones), using ConvNet (TensorFlow team 2016), ResNet (He al. (016), and MobileNetV2 (Howard al. 2017, Sandler et 2018)) FEDBE consistently outperforms FEDAVG, especially when the neural network architecture goes deeper. Moreover, FEDBE can be compatible with existing FL algorithms that regularize clients'
learning or leverage server momentum (Li et al. 2020a Sahu et al.] 2018 Karimiredly et al. 2020 Hsu al., 2019) and further improves upon them. Interestingly, even if the unlabeled server data have a different distribution or domain from the test data (c.g., taken from a different dataset), FEDBE can
still maintain its accuracy, making it highly applicable in practice.
2 RELATED WORK (MORE IN APPENDIX A
Federated learning (FL). In the multi-round setting, FEDAVG (McMahan et al. 2017) is the standard approach. Many works have studied its effectiveness and limitation regarding convergence, robustness, and communication cost, especially in the situations of non-i.i.d. clients. Please see (Appendix for
list of works. Many works proposed to improve FEDAVG. FEDPROX (Li et al. 3020a[]Sahu et al.)
2018), FEDDANE (Li et a 2019). Yao et al. (2019), and SCAFFOLD (Karimireddy et al. (0220) designed better local training strategies to prevent clients' model drifts. /Zha et al. (2018) studied the use of shared data between the server and clients to reduce model drifts. Reddi et al. (2020) and Hsu cet al. (2019) designed better update rules for the global model by server momentum and adaptive optimization. Our FEDBE is complementary to and can be compatible with these efforts.
In terms of model aggregation. Yurochkin et al. (2019) developed a Bayesian non-parametric approach to match clients' weights before average, and FEDMA (Wang al.) 2020) improved upon it by iterative layer-wise matching. One drawback of FEDMA is its linear dependence of computation and communication on the network's depth, not suitable for deeper models. Also, both methods not yet applicable to networks with residual links and batch normalization (Toffe& & Szegedy 2015). We improve aggregation via Bayesian ensemble and knowledge distillation, bypassing weight matching.
Ensemble learning and knowledge distillation. Model ensemble is known to be more robust and accurate than individual base models (Zhoul 2012. Dietterich 2000 [Breiman 1996). Several recent works Anil et al. 2018 Guo et al. [0200] Chen et al. 2020) investigated the use of model ensemble and knowledge distillation (Hinton et al. (015)) in an online fashion to jointly learn multiple models, where the base models and distillation have access to the centralized labeled data or decentralized data of the same distribution. In contrast, client models in Fl. are learned with isolated and likely
DDstiillatinn from the ensemble of clients' models was explored in Guha al. 2019 for one- one-round federated setting. Our work can he viewed as an extension the multi-round setting, by sampling more and higher-quality models as the hases for more robust ensemble.
We note that, the ensemble predictions can he noisy but still more accurate than weight average (see Figure 3 and subsecctoto C.2
2

---

Published as a conference paper at ICLR 2021
non-ii.. and limited data; our distillation is performed without labeled data. We thus propose to sample base models of higher quality for Bayesian ensemble and employ SWA for robust distillation.
Knowledge distillation in FL. Guha et al. (2019) considered one-round FL and applied distillation to obtain a global model from the direct ensemble of clients' models. A similar idea was used in (Papernot et al. 2017) in a different context. Our method can be viewed as extension of (Guha et al. 2019) to multi-round FL., with higher-quality base models being sampled from a global distribution for more robust ensemble. Knowledge distillation was also used in (Li & Wang 2019) and (Jeong et al. 2018) but for different purposes. & Wang (011)) performed ensemble distillation for each client, aiming learn strong personalized models but not the global model. Jeong al.) (2018) aimed to speed up communication by sending averaged logits of clients' data, not models, between clients and the server. The clients then use the aggregated logits to regularize local training via distillation. The accuracy. however, drops drastically compared to FEDAVG in exchange for faster communication.
contrast, we distill on the server using unlabeled data collected at the server, aiming to build stronger global model. The most similar work to ours is concurrent work by Lin et al. 2020 which also employs ensemble distillation on the server in a multi-round setting. Our work is notably different from all the above methods by taking the Bayesian perspective to sample better base models and investigating SWA for distillation, significantly improving the performance on multi-round FL.
3 BAYESIAN MODEL ENSEMBLE FOR FEDERATED LEARNING
3.1 BACKGROUND: FEDAVG
Federated learning (FL) usually involves a server coordinating with many clients to jointly learn global model without data sharing. in which FEDAVG (McMahan et al. 2017) in standard approach. Denote by S the set of clients, D:- = {(x,.. }n)] n=1 the labeled data of client i, and wo the weights of the current global model, FEDAVG starts with client training of all the clients in parallel, initializing each clients' model w; with v and performing SGD for K steps with step size T7
(1) where é is a loss function and B& is the mini-batch sampled from D; at the kith step. After receiving all the clients' models iw E 8). given [D] E, [D,]. FEDAVG performs weight average to update the global model w
Client training:
w; - w. m V wt), for
,K,
Dr D
Model aggregation (by weight average):
w
With the updated global model w. FEDAVG then starts the next round of client training. The whole procedure of FEDAVG therefore iterates between Eqqation and |Equttinn 2 for R rounds.
In the case that Dj is i.i.d. sampled from the aggregated data D = ULEs D., FEDAVG has been shown convergent the ideal model w" learned directly from D in a centralized manner (Stich) 2019, Haddadpour & Mahdavi 2019 Khaled et al. (0200). In reality, however, the server has little control and knowledge about the clients. Each client may have different data distributions in the input (e.g., image distribution) output (e.g., label distribution). Some clients may disconnect at certain rounds. All of these factors suggest the non-i.i.d. nature of federated learning in practice, under which the effectiveness of FEDAYG can largely degrade (Zhao et al. 2018 Li et al. 2020b Hsu et al. 2019). For example, Karimireddy et al. (2020) show that v in Equation can drift away from w".
3.2 BAYESIAN PERSPECTIVE
We propose to view the problem of model drift from Bayesian perspective. In Bayesian learning, it is the posterior distribution p(w)) ((w)D) 'D) of the global model being learned, from which to and w" can be regarded two particular samples (i.e., point estimates). Denote by p(yl æ;; w) the output probability of a global model w, one approach to mitigate model drift is to perform Bayesian inference (Neal 2012 Barber 2012) for prediction, integrating the outputs of all possible models w.r.t. the posterior
p(y|:; [;; D) = p(y) [ w) ) DD)w
(3)
We notice generalized version of it in (He et al. 2020 that improves the computational efficiency.
3

---

Published as a conference paper at ICLR 2021
rather than relying single point estimate. While Equation 3 is intractable in general, we can approximate it by the Monte Carlo method, sampling MZ models for model ensemble
M
Bayesian model ensemble:
D) R
plyla w())) where w (m)) ~ p(w(D)) (4)
M
m=1
The question is: how estimate p(u) (D) in federated learning, given merely client models {w;}
3.3 BAYESIAN MODEL ENSEMBLE WITH APPROXIMATED POSTERIORS
wverage-Gaussian Gaussian (SWAG) (Maddox ctt al., 2019), for estimating the posterior. SWAG employed a cyclical or constant learning rate in SGD, following SWA [Izmailov al. 2018). SWAG then constructs Gaussian distribution p( p(w] 'D) by fitting the parameters to the model weights it traverses in SGD.
We resort recently proposed idea, named stochastic weight average-
In federated learning, by rewriting w, v 9.. where 9: (ww; w denotes the K-step stochastic gradient on a mini-batch D; C 'D (McMahan et al. 2017). we can indeed view each client's
model wP taking K-step SGD to traverse the weight space of global models.
Gaussian. To this end, we propose to fit diagonal Gaussian distribution Ni,, Emng) to the clients' models {un:} following (Maddox et al. 2019).
Dr
D2 wi 4x²
Σ
W
Ediag
diag Σ
(5)
D
Im=1 for model ensemble (cf. Equation 4). Here ())² means taking clement-wise square, We note that, both the clients' models {w,} and FEDAVG w possible samples from /N (Nu, Eding).
from which we can sample (w(m) ~ NN,, Eding) )
Dirichlet. We investigate another way to construct p( DD), inspired by the fact that an averaged stochastic gradient is general closer the true gradient than individual stochastic gradients (Had- dadpour & Mahdavil (2019; Izmailov et al. 2018 Liang et al., 2019 Stich 2019: Zhou & Cong 2017).
By viewing each client's model as w; v gi such fact suggests that convex combination (i.e.. weighted average) of clients" models can lead to a better model than each client alone:
w = Σ Dr
Ep Wo Drl w. = w - Σ
-E
Ex 7h (D, 9..
7: D
(6)
where = m.... Ms] E A is vector on the (IS) - 1) Dismlle.. To this end, we use Dirichlet distribution Dir(a)) to model the distribution of y, from which we can then sample w by
w (m)
(m) 7
[D,
((m) ~ p(r) = p(n, mssi) = B(a)
1
II
0 (7)
wi
Ee
(m) Dyyl
where [a1.... a si] 0 is the parameter of Dirichlet distribution, and B(a) is the multivariate beta function for normalization. We study different a in subsecctio C.T
To sum up, Bayesian model ensemble in federated learning takes the following two steps:
Construct p(w)D) (D) from the clients' models {w;} } (cf. Equation 5 or Equation
Sample (m( (m) ~ (D) and perform ensemble (cf. Equation 4)
Analysis. We validate Bayesian model ensemble with three-class classification problem on the Swiss roll data in Figure (a). We consider three clients with the same amount of training data: each has 80% data from one class and 20% from the other two classes, essentially a non-i.i.d. case. We apply FEDAVG to train a two-layer MLP for 10 rounds (each round with 2 epochs). We then show the test accuracy of models sampled from [Equation 7 (with a 0.5 x 1) - the corners of the triangle (i.e., A in Figure (b) correspond to the clients: the position inside the triangle corresponds to the coefficients. We see that, the sampled models within the triangle usually have higher accuracy than the clients' models. Surprisingly, the best performing model that be sampled from Dirichlet distribution is not FEDAVG (the center of the triangle), but the one drifting to the bottom right. This suggests that Bayesian model ensemble can lead to higher accuracy (by averaging over sampled
4

---

Published as a conference paper at ICLR 2021
Text Data
Clont
6200
Correct Incorrect
Clients
6575
Fedlig
5.00
6025
Clent
Clent II
6080
Samples
6355
6:50
633
20
(a)
(b)
Figure I: An illustration of models that can be sampled from Dirichlet distribution (Equation 7 (a) three-class toy data with three clients, each has non-i.i.d. imbalanced data. (b) We show the sampled model's corresponding Y (position in the triangle) and its test accuracy (color). FEDAVG is at the center; clients' models are at corners. The best performing model (star) is not at the center, drifting away from FEDAVG. (c) Histograms of (in)correctly predicted examples at different confidences (x-axis) by sampled models and clients.
models) than FEDAVG alone. Indeed, by combining 10 randomly sampled models via Equation 4 Bayesian model ensemble attains a 69% test accuracy, higher than 64% by FEDAVG. Figure 1 (c)
further shows that the sampled models have a better alignment between the prediction confidence and accuracy than the clients' (mean results of clients or samples). See subsection C. for details.
To further compare FEDAVG and ensemble, we linearize p(y [x; -) at (Izmailov et al. 2018
plylæ; 2; w (m)) = p(u x; w) + (Vpíylæ; w)) Ω m/m)) + ((1(((m))))
(8)
where (m) (m) vn and (..)) ) is the dot product. By averaging the sampled models, we arrive at p(yl an; w (m)) p æz; u)) (Vrlyl w) M
M Σ
((())) 0( 0(2²),
where Ω = maxm || In federated learning, especially in the non-i.i.d. cases, can be quite large. Bayesian ensemble thus can have a notable difference (improvement) compared FEDAVG.
FEDBE
Algorithm 1: FEDBE (Federated Bayesian Ensemble)
Server input :initial global model w. SWA scheduler 7/SWA, unlabeled data 24 = sez)==1
Bayesian model ensemble, however, can- not directly benefit multi-round federated learning, in which a single global model must sent back to the clients to con- tinue client training. We must translate the prediction rule of Bayesian model en- semble into single global model.
Client i's input: step size rp, local labeled data D,;
for + 1 to R do
Sample clients S C (1, N); Communicate for each client i E S in parallel do Initialize local model w; <- w:
all clients E S;
w. <- Client training (w, D)): Equation
To this end, we make an assumption that can access set of unlabeled data u = [ at the server. This can eas- ily be satisfied since collecting unlabeled data is simpler than labeled ones. We use for two purposes. On one hand, we use U to memorize the prediction rule of Bayesian model ensemble, turning into pseudo-labeled set T {( P,)) /
Communicate w, to the server; end
Construct to
Eies Eves PD w:; Construct global model distribution p(w) (D) from
P
{wa i E S);
Equation or Equation 7
Sample M global models
(m) p(w D) mm=
Construct (mn((m)) )
{10) [wn;; i E S) J (m)) Jm=
Construct T (I,, Dy H where
where p. 37 M P(yl ay: (m)) is probability vector. On the other hand, we use 7 as supervision to train a global model w, aiming to mimic the prediction rule of Bayesian model ensemble on 4.
p:j X7 Σ ( to ));
Equation 4
Knowledge distillation: < SWA(w 75Ww))
end Server output :w.
This process is reminiscent of knowledge distillation (Hinton et al., 2015) to transfer knowledge from a teacher model (in our case, the Bayesian model ensemble) to student model (a single global model). Here we apply a cross entropy loss to learn w - Σ, DJ log(p(y) (p(p(yla,,; w)).
SWA for knowledge distillation. Optimizing w using standard SGD, however, may arrive at a suboptimal solution: the resulting can have much worse test accuracy than ensemble. We identify
5

---

Published as a conference paper at ICLR 2021
one major reason: the ensemble prediction py can be noisy (e.g., arg max, Pj[c] is not the true label of æ,), especially in the early rounds of FL. The student model w thus may over-fit the noise. We note that, this finding does not contradict our observations in susssecctionn 3.3; Bayesian model ensemble has higher test accuracy than FEDAVG but is still far from being perfect (i.e., 100% accuracy). To address this issue, we apply SWA (Izmailov al. 2018) to train w. SWA employs cyclical learning rate schedule in SGD by periodically imposing sharp increase in step sizes and averages the weights of models traverses, enabling w jump out of noisy local minimums. As will be shown in (escctttinn 5) SWA consistently outperforms SGD in distilling the ensemble predictions into the global model.
We name our algorithm FEDBE (Federated Bayesian Ensemble) and summarize it in algorithm While knowledge distillation needs extra computation at the server, it is hardly a concern as the server likely computationally rich. (See (Seeeccuse D.I] for details.) We also empirically show that a small number of sampled models (e.g.. M = 10 ~ 20) are already sufficient for FEDBE to be effective.
5
EXPERIMENT
5.1 SETUP (MORE DETAILS IN APPENDIX B
Datasets, models, and settings. We use CIFAR- (Krizhevsky et al. 2009), both contain 50K training and 10K test images, from 10 and 100 classes. We also use Tiny-ImageNet (Le & Yang 2015), which has 500 training and 50 test images per class for 200 classes. We follow (MeMahan et al. 2017) to use a ConvNet (LeCun et al.) 1998) with 3 convolutional and 2 fully-connected layers. We also use ResNet- 32, 44, 56) (He et al. 2016) and MobileNetV2 (Sandler et al. 2018). We split part of the training data to the server as the unlabeled data, distribute the rest to the clients, and evaluate on the test set. We report mean + standard deviation (std) over five times of experiments.
2020 Li et al. Implementation details. As mentioned in (McMahan et al. 2017 Wang et al. 2020b), FEDAVG is sensitive to the local training epochs E per round (E K
in Equation
D..
Thus, in each experiment, we first tune E from [1, 5, 10, 20, 30, 40) for FEDAVG and adopt the same E to FEDBE. Li et al. (2020b); Reddi et al. 2020) suggested that the local step size tii (see 1 must decay along the communication rounds in non-i.i.d. settings for convergence. We set the iniiial 7): as 0.01 and decay it by 0.1 at 30% and 60% of total rounds, respectively. Within each round of local training, we use SGD optimizer with weight decay and a 0.9 momentum and impose no further decay local step sizes. Weight decay is crucial in local training (cf. /subsection B.3) For ResNet and MobileNetV2, we use batch normalization (BN), See Seessubsettion C.5 for discussion on using group normalization (GN) (Wu& & He 2018 Hsich al. 2020). which converges much slower.
Equation
Baselines. Besides FEDAVG, we compare to one-round training with 200 local epochs followed by model ensemble the end (1-Ensemble). We also compare to vanilla knowledge distillation (v-Distillation), which performs ensemble directly over clients' models and uses a SGD momentum optimizer (with batch size of 128 for 20 epochs) for distillation in each round. For fast convergence, we initialize distillation with the weight average of clients" models and sharpen the pseudo label as Py[c] < P [c] E.' p. similar to (Berthelot et al., 2019). We note that, v-Distillation highly similar to (Lin et al. 2020) except for different hyper- parameters. We also compare FEDPROX (Li ct al. 2020a) and FEDAVOM (Hsu et al. 2019) on better local training and using server momentum.
FEDBE, We focus Gaussian (cf. Equation 5). Results with Dirichlet distributions are in sub section C.I. We sample M=10 models and combine them with the weight average of clients and
individual clients for ensemble. For distillation, we apply SWA (zmaillv al. 2018), which uses cyclical schedule with the step size JJSWA decaying from 1e-3 to 4e-4, and collect models at the end of every cycle (every 25 steps) after the 250th step. We follow other settings of v-Distillation (c.g., distill for 20 epochs per round). We average the collected models to obtain the global model.
5.2
MAIN STUDIES: CIFAR-10 10 WITH NON-LLD. CLIENTS USING DEEP NEURAL NETWORKS
Setup. We focus on CIFAR-10. 10. We randomly split 10K training images to be the unlabeled data at the server. We distribute the remaining images to 10 clients with two non-i.i.d cases. Step: Each client has 8 minor classes with 10 images per class, and 2 major classes with 1,960 images per class, inspired by (Cao et al. 2019). Dirichlet: We follow (Hsu et al. 2019) to simulate a heterogeneous

---

Published as a conference paper at ICLR 2021
Table Mean=std of test accuracy (%) on non-i.id. CIFAR-I0 10. x: trained with 50K images without splitting.
Non-i.id. Type
Method
ConvNet
ResNer20
ResNet32
ResNet44
ResNet56
I-Ensemble
60.5:10.28
49.9.10.46
35.51:.50
32.8.10.8
23.3:10:52
FEDAVG
72.0=0.25
70.2=t017
66.5=0.36
60.5:0:26
51,4+0.15
Step
v-Distillation
69.2:t0.18
72.6=0062
68.4=0.33
60.4=0.53
56,4+1110
FEDBE (w/o SWA)
7..1=1.21
74.9=1141
71.1=0.75
61.0=0.75
56.6+0.85
FEDBE
74.5:10.51
77.5:10.22
72.7=0.27
65.5
60.7
1-Ensemble
63.310.56
1.02:
39.5=0.78
31.5+10.77
27.2:t0.65
FEDAVG
72.3+0012
04.4
73.4+0.23
67.1+0.54
62.2++0.45
Dirichlet
y-Distillation
67.7+0098
73.1
70.8+0.64
66.9++0.5
62.8++0.66
FEDBE (w/o SWA)
70.1+0042
75.9+
73.9=0.55
68.2+0072
63.2+0071
FEDBE
73.9:10.45
7.32
77.7=0.45
71.5+0..8
67.0
Centralizedk
SGD
84.5
91.7
92.6
93.1
93.4
Table 2: Compatibility of FEDBE with FEDAvGM and FEDPROX non-i.id. CIFAR-10.
Non-i.id. Type
Method
ConvNet
ResNer20
ResNet32
ResNer44
ResNe156
FEDPROX
72.5+0.71
71.1+0.52
67.71:026
60.410.71
54.9=2066
FEDBE +FEDPROX
74.9+10.88
77.7::0.45
72.9-0044
64.5+0.37
60.1=0.62
Step
FEDAYGM
72.3++0.55
73.2+0057
70.0+0062
59.9+00.55
52.7+-0.99
FEDBE +FRDAVGM
74.5+0047
78.0+
73.6+-0.50
65.5++0.40
59.7+0.51
FEDPROX
72.6+00.8
76.1+0049
73.4+0051
68. 0..19
60.9++0.6
FEDBE +FEDPROX
74.6+-0.35
78.7+0049
77.3+:0.00
71.7++0.43
66.5=00.41
Dirichlet
FEDAvGM
73.0++0.4
76.5+0.44
75.5+
67.7+:0.6
58.9=0072
FEDBE +FEDAVGM
74.410.49
78.510066
78.5
72.0-0.51
67.0-10.55
partition for N clients on C classes. For class c, we draw N-dim vector q. from Dir(0.1) and assign data to client n proportionally to qe [m]. The clients have different numbers of total images.
Results. We implement all methods with 40 rounds"] except for the one-round Ensemble. We assume that all clients are connected every round. We set the local batch size as 40. Table I| summarizes the results. FEDBE outperforms the baselines by notable margin. Compared to FEDAVG, FEDBE consistently leads to a 9% gain, which becomes larger as the network goes deeper. By comparing FEDBE to FEDBE (w/o SWA) and v-Distillation, for distillation and Bayesian ensemble with sampled models. We note that, FEDAVG outperforms 1 1-Ensemble and is a par with v- v-Distillationf) justifying (a) the importance of multi-round training; (b) the challenge of ensemble distillation. Please see (subsection f.r an insightful analysis.
see the consistent improvement by SWA
Compatibility with existing efforts. Our improvement in model aggregation is compatible with recent efforts in better local training (Li et al. 2020a; Karimireddy et al. 2020 and using server
momentum (Reddi al. 2020 Hsu al. 2019). Specifically, Reddi et al. 2020) Hsu al. (0099)
applied the server momentum to FEDAVG by treating FEDAVG in each round as step of adaptive
optimization. FEDBE can incorporate this idea by initializing distillation with their FEDAVG. Table
shows the results of FEDPROX (Li et al.) 2020a) and FEDAVGM (Hsu et al. 2019). w/ or w/o FEDBE.
FEDBE can largely improve them. The combination even outperforms FEDBE alone many cases.
Effects of Bayesian Ensemble. We focus on the Step setting. We compare different combi- nations of client models C: [w, client weight average A: w, and M samples from Gaussian S: (m) mm to construct the distillation tar- gets T for FEDBE in algorithh I As shown in Table 3 sampling global models for Bayesian ensemble improves the accuracy. Sampling M 10 ~ 20 samples (plus weight average and clients to form ensemble) is sufficient to make FEDBE effective (see Figure
Table 3: FEDBE distillation targets. A: client average; C: clients; S: samples. Distillation Targets
ConvNet ResNer20
72.6110.88 73.4.0..66 73.1+-0.46
75.2++0.11
S+A A+C
73.9+0.33 +0.47
73.0+-0.36
75.440.38
S+C
74.0.10.66
77.9
S+A+C
74.510.51 77.5
Ground-truth labels 76.6.10.21 80.2:1023
WWe observe very little gain after 40 rounds: adding 60 rounds only improves FEDAVG (ConvNet) by 0.7%.
In contrast Lin et al. (2020), we add weight decay to local client training, effectively improving FEDAVG.
7

---

Published as a conference paper at ICLR 2021
78
76
12.5
Test
70.0
TacDC- GT
Rochet20 Conoter
sendag
65.0
10
20
30
40
#Samples
y Com layers Figure 4: # of layers (ConvNet).. GT: with ground -ruut targets.
Figure 2: # sampled models in
Figure 3: FEDAVG while monitor- ing the Bayesian ensemble.
FEDBE.
Table 4: FEDBE on non-i.i.d CIFAR- 10 with different unlabeled data 24
Non-i.id. Type
u
[
ConvNet
ResNet20
ResNet32
ResNet44 ResNet56
CIFAR-10
10K
74.5++0.51
77.5+0.42
72.7+0.27
65.5+0.32 60.7+00.45
Step
CIFAR- 100
50K
74,4+0.45
78.2++0.58
02.2+0.35
+0.37
61.0+0049
Tiny- Imageeege
100K
74.5+0.64
77.1=0.51
72.3+0.43
64.5#0.51
60.9+0.32
CIFAR-10
10K
13.99+0.45
78.2++0.36
77.7+0.45
71.5+0.38 67.0++0.0
Dirichlet
CIFAR- 100
50K
73.5+0.41
78.6+0.63
76.5+0.61
72.0+0.71
66.9+0.57
Tiny- Iny-emmeeNet
100K
74.0++0.35 78.2+0.72
76.7=0.52
71.6 67.3++.32
78
78
Test Accurecs
Teat Accuracy
70-
74
ResNet20
FedBE ConwNir
-5
FecBi ReaNet50 Commlet
Test
72
Tectt Realee20
65
RedAvg-nestaat20
FoomE
Cenvueg-onvttet
70
Commlett
3k
5k
10c
15k
25c
68
24
1/
68
24c
12.
20c
53k
AServer data
dServerr diata
derreer datta
(a) CIFAR-10.
(b) CIFAR-100.
(c) Tiny-ImageNet.
Figure 5: Effects on varying the size and domains of the server dataset on CIFAR- 10 experiments.
Bayesian model ensemble vs. weight average for prediction. In Figure 3 we perform FEDAVG and show the test accuracy at every round, together with the accuracy by Bayesian model ensemble, on the Step-non-i..d CIFAR-10 experiment using ResNet20. That is, we take the clients' models learned with FEDAVG to construct the distribution, sample models from it, and perform ensemble for the predictions. Bayesian model ensemble outperforms weight average at nearly all the rounds, even though it is noisy (i.e., not with 100% accuracy).
Effeets of unlabeled data. FEDBE utilizes unlabeled data to enable knowledge distillation. Figure 5a studies the effect of [T] we redo the same Step experiments but keep 25K training images away from clients and vary [44] in the server. FEDBE outperforms FEDAVG even with 1K unlabeled dataset (merely 4% of the total client data). We note that, FEDAVG (ResNet20) trained with the full 50K images only reaches 72.5%, worse than FEDBE, justifying that the gain by FEDBE is not simply from seeing more data. Adding more unlabeled data consistently but slightly improve FEDBE.
We further investigate the situation that the unlabeled data come from a different domain or task. This simulate the cases that (a) the server has little knowledge about clients' data and (b) the server cannot collect unlabeled data that accurately reflect the test data. In Table 4 we replace the unlabeled data to CIFAR- 100 and Tiny-Imagen ImageNet. The accuracy matches or even outperforms using CIFAR- 10, suggesting that out-of -o-mann unlabeled data are sufficient for FEDBE. The results also verify that FEDBE uses unlabeled data mainly as a medium for distillation, not a peep at future test data.
In Figure 5bd Figure 5c we investigate different sizes of CIFAR-100 or Tiny-ImageNet as the unlabeled data (cf. Table 4b. We found that even with merely 2K unlabeled data, which is 5% of the total 40K CIFAR- 10 labeled data and 2 ~ 4% of the original 100K unlabeled data, FEDBE can already outperform FedAvg by a margin. This finding is aligned with what we have included in Figure 5a where showed that small amount of unlabeled data is sufficient for FEDBE to be effective. Adding more unlabeled data can improve the accuracy but the gain is diminishing.
Network depth. Unlike in centralized training that deeper models usually lead to higher accuracy (bottom row in Table I trained with 200 epochs), we observe an opposite trend in FL: all methods suffer accuracy drop when ResNets go deeper. This can be attributed to (a) local training over-fitting
8

---

Published as a conference paper at ICLR 2021
Table 5: Partial participation (Tiny-ImageNe)
Table 6: Systems heterogeneity (non-i.i.d. CIFAR-10)
Method ReesNet20 MobileNerV2
Method
ConvNet
ResNer20 ResNer32
i.id
FEDAVG 332.4+0.68
=0.98
FEDAVG 70.6+0.46 69. 69.9+00.59 64,0+0.50 FEDPROX 71.2:10.55 69.4110.48 65.9.10.3
FEDBE 33..41.0.58
28.9-1.15
FEDAVG 027.5+0.78
25.5-11.23
FEDBE 73.3++0.56
77.1+0.61 70.2++0.9 +FEDPROX 73.7.10.24 77.510..1 71.610..37
non-i.id
FEDBE 32.4.10.81
27.8-0099
to small and non-i.i.d. data or (b) local models drifting away from each other. FEDBE suffers the least among all methods, suggesting it as promising direction to resolve the problem. To understand the current limit, we conduct study in Figure 4|by injecting more convolutional layers into ConvNet (Step setting). FEDAVG again degrades rapidly, while FEDBE is more robust. If we replace Bayesian ensemble by the CIFAR-10 10 ground truth labels as the distillation target, FEDBE improves with more layers added, suggesting that how to distill with noisy labeled targets is the key to improve FEDBE.
5.3 PRACTICAL FEDERATED SYSTEMS
Partial participation. We examine FEDBE in a more practical environment: (a) more clients are involved, (b) each client has fewer data, and (c) not all clients participate in every round. We consider a setting with 100 clients, in which 10 clients are randomly sampled at each round and iterates for 100 rounds, similar to (McMahan et al., 2017). We study both i.i.d. and non-i.i.d (Step) cases on Tiny-ImageNet, split 10K training images to the server, and distribute the rest to the clients. For the non-i.i.d case, each client has 2 major classes (351 images each) and 198 minor classes (1 image each). FEDBE outperforms FEDAVG (see Table 5). See(subsection C.7 for results on CIFAR-100.
Systems heterogeneity. In real-world FL, clients may have different computation resources, leading to systems heterogeneity (Li et al. 2020a). Specifically, some clients may not complete local training upon the time of scheduled aggregation, which might hurt the overall aggregation performance. We follow (Li et al. 2020a) to simulate the situation by assigning each client local training epoch E, sampled uniformly from (0, 20), and aggregate their partially-trained models Table 6 summarizes the results on non-i.i.d (Step) 10. FEDBE outperforms FEDAVG and FEDPROX (Li et al. 2020a).
DISCUSSION
Privacy. Federated learning offers data privacy since the server has no access to clients' data. It is worth noting that having unlabeled data not collected from the clients does not weaken the privacy if the server is benign, which is a general premise in the federated setting (McMahan et al. 2017). For instance, if the collected data are de-identified and the server does not intend to match them to clients, clients' privacy is preserved. In contrast, if the server is adversarial and tries infer clients' information, federated learning can be vulnerable even without the unlabeled data: e.g., federated learning may not satisfy the requirement of differential privacy or robustness membership attacks.
Unlabeled data. Our assumption that the server has data is valid in many cases: car company may collect its own data but also collaborate with customers to improve the system. Bassily et al. [2020a] b) also showed real cases where public data is available in differential privacy.
e.g., self-driving
7 CONCLUSION
Weight average in model aggregation is one major barrier that limits the applicability of federated learning to i.i.d. conditions and simple neural network architectures. We address the issue by using Bayesian model ensemble for model aggregation, enjoying a much robust prediction at very low cost of collecting unlabeled data. With the proposed FEDBE, we demonstrate the applicability of federated learning to deeper networks (i.e., ResNet20) and many challenging conditions.
ACKNOWLEDGMENTS
We are thankful for the generous support of computational resources by Ohio Supercomputer Center and AWS Cloud Credits for Research.

---

Published as a conference paper at ICLR 2021
REFERENCES
Rohan Anil, Gabriel Pereyra, Alexandre Tachard Passos, Robert Ormandi, George Dahl, and Geoffrey Hinton. Large scale distributed neural network training through online distillation. 2018. URL
ew net teee/pppfia=rcktueeeee
David Barber. Bayesian reasoning and machine learning. Cambridge University Press, 2012.
Raef Bassily, Albert Cheu, Shay Moran, Aleksandar Nikolov, Jonathan Ullman, and Zhiwei Steven Wu. Private query release assisted by public data. In ICMI., 2020a.
Raef Bassily, Shay Moran, and Anupama Nandi. Learning from mixtures of private and public populations, arXiv preprini arXiv:2008.0033 2020b.
Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman Vaughan. A theory of learning from different domains. Machine learning, 79(1-2)1151-175, 2010.
David Berthelot, Nicholas Carlini, Ian Goodfellow, Nicolas Papernot, Avital Oliver, and Colin A Raffel. Mixmatch: A holistic approach to semi-supervised learning. In NeurIPS, 2019.
Keith Bonawitz, Hubert Eichner, Wolfgang Grieskamp, Dzmitry Huba, Alex Ingerman, Vladimir Ivanov, Chloe Kiddon, Jakub Konecny, Stefano Mazzoechi, H Brendan McMahan, al. Towards federated learning at scale: System design. arXiv preprint arXi::1902.01044 2019.
Leo Breiman. Bagging predictors. Machine learning, 24(2): 140, 1996.
Eric Brochu, Vlad M Cora, and Nando De Freitas. A tutorial on bayesian optimization of expensive cost functions, with application to active user modeling and hicrarchical reinforcement learning. arXiv preprint arXiv: 1012,259,, 2010.
Kaidi Cao, Colin Wei, Adrien Gaidon, Nikos Arechiga, and Tengyu Ma. Learning imbalanced datasets with label-distribution-awaas margin loss. In Advances in Neural Information Processing Systems, pp. 1565-1576, 2019.
Defang Chen, Jian-Ping Mei, Can Wang, Yan Feng, and Chun Chen. Online knowledge distillation with diverse peers. In AAAI, 2020.
Thomas G Dietterich. Ensemble methods in machine learning. In International workshop on multiple classifier systems, pp. 1-15, 2000.
Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht. Essentially no barriers in neural network energy landscape. In ICML, 2018.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversarial training of neural networks. The Journal of Machine Learning Research, 17(1):2096-22) 20200000,,,,,, 2016.
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew G Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. In NeurIPS, 2018.
Boqing Gong, Yuan Shi, Fei Sha, and Kristen Grauman. Geodesic flow kernel for unsupervised domain adaptation. In CVPR, 2012.
Boqing Gong, Kristen Grauman, and Fei Sha. Learning kernels for unsupervised domain adaptation with applications to visual object recognition. IJCV, 109(1-2):3-27, 2014.
Yves Grandvalet and Yoshua Bengio. Semi-supervised learning by entropy minimization. In NIPS. 2005.
Neel Guha, Ameet Talwikar, and Virginia Smith. One-shot federated learning, arXiv preprint /902. / 175, 2019.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. 2017.
10

---

Published as a conference paper at ICLR 2021
Qiushan Guo, Xinjiang Wang, Yichao Wu, Zhipeng Yu, Ding Liang, Xiaolin Hu, and Ping Luo.
Online knowledge distillation via collaborative learning. In CVPR, 2020.
Farzin Haddadpour and Mehrdad Mahdavi. On the convergence of local descent methods in federated learning. arXiv preprint arXiv:19100 14425, 2019.
Chaoyang He, Murali Annavaram, and Salman Avestimehr. Group knowledge transfer: Federated learning of large cnns at the edge. Advances in Neural Information Processing Systems, 33, 2020.
Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. CVPR, 2016.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge neural network. arXiv preprint arXiv:1503.02533 2015.
Andrew G Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam. Mobilenets: Efficient convolutional neural networks for mobile vision applications. 2017.
Kevin Hsieh, Amar Phanishayee, Onur Mutlu, and Phillip B Gibbons. The non-iid data quagmire of decentralized machine learning. In ICML, 2020.
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.06335. 2019.
Gao Huang, Yixuan Li, Geoff Pleiss, Zhuang Liu, John E Hoperoft, and Kilian Q Weinberger. Snapshot ensembles: Train get for free. In ICLR, 2017.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv:1502.03160 2015.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. In UAI, 2018
Eunjeong Jeong, Seungeun Oh, Hyesung Kim, Jihong Park, Mehdi Bennis, and Seong-Lyu Kim. Conmunication-efficceenn on-device machine learning: Federated distillation and augmentation under non-iid private data. arXiv preprint arXiv:1811.1142 2018.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for on-device federated learning. ICMI., 2020.
Khaled, K Mishchenko, and P Richtárik. Tighter theory for local sgd on identical and heterogeneous data. In ALSTATS, 2020.
Durk P Kingma, Shakir Mohamed, Danilo Jimenez Rezende, and Max Welling, Semi-supervised learning with deep generative models. In NIPS, 2014.
Jakub Konečný, Brendan McMahan, Felix X Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon. Federated learning: Strategies for improving communication efficiency. arXiv preprint arXiv:1610.05492 2016.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Ya Le and Xuan Yang. Tiny imagenet visual recognition challenge. CS 231N, 2015.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, (6(11):2278-2322 1998.
Daliang Li and Junpu Wang. Fedmd: Heterogenous federated learning via model distillation. arXiv
preprint 10X0:19110003555 2019.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smithy. Feddane: A federated newton-type method. In 2019 53rd Asilomar Conference on Signals, Systems, and Computers, pp. 1227-1231. IEFF, 2019.
11

---

Published as a conference paper at ICLR 2021
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith.
Federated optimization in heterogeneous networks. In MLSys, 2020a.
Xiang Li, Kaixuan Huang, Wenhao Yang. Shusen Wang, and Zhihua Zhang, On the convergence of
fedavg on non-iid data. In ICLR, 2020b.
Xianfeng Liang, Shuheng Shen, Jingchang Liu, Zhen Pan, Enhong Chen, and Yifei Cheng. Variance reduced local sgd with lower communication complexity. arXiv preprint 1XXi:1912.11244 2019.
Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. arXiv preprint arXiv:2006.0722 2020.
Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. JMLR, 9NNov):5579-2605 2008.
Wesley J Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson. A simple baseline for bayesian uncertainty in deep learning. In NeurIPS, 2019.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, et al. Communication--fficcennt learning of deep networks from decentralized data. In AISTATS, 2017.
Radford M Neal. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 2012.
Nicolas Papernot, Martín Abadi, Ulfar Erlingsson, Ian Goodfellow, and Kunal Talwar. Semi- supervised knowledge transfer for deep learning from private training data. ICLR, 2017.
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. arXiv preprint arXiv:2003.00295, 2020.
Amirhossein Reisizadeh, Aryan Mokhtari, Hamed Hassani, Ali Jadbabaie, and Ramtin Pedarsani. Fedpaq: A communication-effiieent federated learning method with periodic averaging and quanti- zation. arXiv preprint 10X1v:1909.13014 2019.
Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. In ICLR, 2020.
Anit Kumar Sahu, Tian Li, Maziar Sanjabi, Manzil Zaheer, Ameet Talwalkar, and Virginia Smith. On the convergence of federated optimization in heterogeneous networks. arXiv preprint arXiv:1812..01127, 2018.
Kuniaki Saito, Kohei Watanabe, Yoshitaka Ushiku, and Tatsuya Harada. Maximum classifier discrepancy for unsupervised domain adaptation. In CVPR, 2018.
Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen. Mo- bilenetv2: Inverted residuals and linear bottlenecks. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 4510-4520, 2018.
Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet S Talwalkar. Federated multi-task learning. In NIPS, 2017.
Sebastian Stich. Local sgd converges fast and communicates little. In ICLR, 2019.
Antti Tarvainen and Harri Valpola. Mean teachers are better role models: Weight-average consis- tency targets improve semi-supervised deep learning results. In Advances in neural information processing systems, pp. 1195-1204, 2017.
TensorFlow team. Tensorflow convolutional neural networks tutorial. tensorflow. wwtp://tttttootta//eeeeeeee.eeee
http:/ www
2016.
Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Federated learning with matched averaging. In ICLR, 2020.
12

---

Published as a conference paper at ICLR 2021
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference computer vision (ECCV), pp. 3-19, 2018.
Qiang Yang, Yang Liu, Tianjian Chen, and Yongxin Tong. Federated machine learning: Concept and applications. ACM Transactions Intelligent Systems and Technology (TIST), 10(2):1-19, 2019.
Xin Yao, Tianchi Huang, Rui-Xiao Zhang, Ruiyu Li, and Lifeng Sun. Federated learning with unbiased gradient aggregation and controllable meta updating. arXiv preprint arXiv: 1910.08234, 2019.
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Trong Nghia Hoang, and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. In ICML., 2019.
Zhengming Zhang, Zhewei Yao, Yaoqing Yang. Yujun Yan, Joseph E Gonzalez, and Michael W Mahoney, Benchmarking semi-supervised federated learning. arXiv preprint arXiv:2008. 11008.11364, 2020.
Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data. arXiv preprint arXiv:/806.0000 2018.
Fan Zhou and Guojing Cong. On the convergence properties of k-step averaging stochastic gradient descent algorithm for nonconvex optimization. arXiv preprint 17Xivv1708.01012, 2017.
Zhi-Hua Zhou. Ensemble methods: foundations and algorithms. CRC press, 2012.
Xiaojin Jerry Zhu. Semi-supervised learning literature survey. Technical report, University of Wisconsin-Madison Department of Computer Sciences, 2005.
Martin Zinkevich, Markus Weimer, Lihong Li, and Alex J Smola. Parallelized stochastic gradient descent. In NIPS, 2010.
13

---

Published as a conference paper at ICLR 2021
SUPPLEMENTARY MATERIAL
We provide details omitted in the main paper.
Appendix A
additional related work (cf. [sccionn of the main paper).
Appendix B
details of experimental setups (cf. (ubseection 5.1 of
of the main paper).
Appendix C
additional experimental results and analysis (cf. subsection 5.2 of the main
paper).
Appendix
additional discussions (cf. Section 4 and subsection 5.2 of the main paper).
Appendix
additional analysis to address reviewers' comments.
A ADDITIONAL RELATED WORK
Federated leatning (FL). In the multi-round federated setting, FEDAVG (McMahan et al. 2017> is the standard approach. Many works have studied its effectiveness and limitation regarding convergence (Khaled et al. 2020 Li al. | Karimireddy al. 2020 Li et al. 20200b Liang et al. 2019 Stich 2019; Zhao et al. 201 8, /Zhou & Congl 2017 Haddadpour & Mahdavi 2019).
system robustness (Li et al. 2020a Smith et al. 2017] Bonawitz al. | (001)) and communication
cost (Konečný et al. 2016 Reisizadeh et al. 2019), especially for the situations that clients are not
i..dd et al. 2020b| Zhao et al. 2018 Li et al. 2020a; Sahu et al. 2018) and have different data distributions, stability, etc.
Ensemble learning and stochastic weight average. Recent works (Huang al. 2017 Draxler et al. 2018 Garipov et al. 2018) have developed efficient ways to obtain the base models for ensemble; c.g., by employing a dedicated learning rate schedule to sample models along a single pass of SGD training (Hsu et al., 2019). SWA (Maddox et al. 2019) applied the same learning rate schedule but simply took weight average over the base models to obtain a single strong model. We apply S SWA, but for the purpose of learning with noisy labels in knowledge distillation.
Bayesian deep learning. Bayesian approaches (Neal 2012 Barber 2012 Brochu et al. (010) incor- porate uncertainty in decision making by placing distribution over model weights and marginalizing these models to form a whole predictive distribution. Our work is inspired by (Maddox et al. 2019 which constructs the distribution by fitting the parameters traversed models along SGD training.
Others. Our work is also related to semi-supervised learning (SSL) and unsupervised domain adaptation (UDA). SSL leverages unlabeled data to train better model when limited labeled data
provided (Grandvalet & Bengio 2 Kingma et al. 2014 Tarvainen & Valpola 2017] Berthelot al., 2019: Zhu 2005); UDA leverages unlabeled data adapt model trained from a source domain to different but related target domain (Gong et al. 2012 2014; Ganin et al. 2016 Saito jet al. 2018 Ben-David et al. 2010). We also leverage unlabeled data, but for model aggregation. We note that, UDA and SLL generally assume the access to labeled (source) data, which is not the case in federated learning: the server cannot access clients' labeled data.
B EXPERIMENTAL SETUPS
B. IMPLEMENTATION DETAILS
As mentioned in the main paper (subsection 5.1). we select the number of local epochs, used in training client model within one round of communication, according to the performance of FEDAVG. Other algorithms, like FEDBE and v- v-Distillation, then follow the same numbers. For ConvNet and ResNet experiments, we fixed 1-Ensemble baselines, tuned E from [10, 20, 200).
20. We use E = 10 for MobileNet experiments. For all
We observed that applying weight decay in local client training improves all FL methods, but the suitable hyper-paue parem-prrameter can be different for different methods different network architectures. Tuning it specifically for each method is thus essential for fair comparisons. We search the weight decay hyper-parameter for each network and each method in [1e-3, 1e-4] with validation set.
14

---

Published as a conference paper at ICLR 2021
For methods with distillation (FEDBE and v-Distillation), We tune the epochs for v-Distillation from [1, 5, 10, 20, 30, 40) and find 20 to be stable across different setups. We apply 20 to FEDBE as well.
In constructing the pscudo-labeled data T, perform inference the unlabeled data U in the server without data augmentation. We perform data augmentation in both local training on D; and knowledge distillation on TT The 32 x 32 CIFAR images are padded 2 pixels each side, randomly flipped horizontally, and then randomly cropped back to 32 x 32. The 64 x 64 Tiny-ImageNet images are padded pixels each side, randomly flipped horizontally, and then randomly cropped back to 64 x 64. In Table 4 of the main paper, we resize images of Tiny-ImageNer to 32 x 32.
For neural networks that contain batch normalization layers (Toffe & Szegedy 2015), we apply the same way as in section 3 construct the global distribution for the layers and we observe no issues in our experiments. SWAG (Maddox et 2019) also reported that performs stably even on very deep networks.
B.2 TRAINING FEDAVG
For the local learning rate ?], we observed that an appropriate value of is important when training on the non-i.i.d local data. The local models cannot converge if the learning rate is set to too large value (also shown in (Reddi et al. (0200), and the models cannot reach satisfying performance within the local epochs with a too-small value as shown Figure 6 Also, unlike the common practice of training a neural network with learning rate decay in the centralized setting, we observed that applying it within each round local training (decay by 0.99 every step) results in much worse client models. The local models would need a large enough learning rate converge with fixed E and we use 0.01 as the base learning rate for local training in all our experiments.
Although decaying nn within each round of local training is not helpful, decaying yp along the rounds of communication could improve the performance. [Wang et al. (2020) and Reddi et al.] (2020) provided both theoretical and empirical studies suggesting that the local learning rate must decay along the communication rounds to alleviate client shifts in non-i.i.d setting. In our experiments, the rth round of communication. the local client training starts with a learning rate m7 which 0.01 if < O.3R, 0.001 if 0.3R < r < 0.6R, and 0.0001 otherwise, where R is the total rounds of communication. In Figure 7 we examined this schedule with different degrees of in the Dirichlet-non-i.i setting. We observed consistent improvements and applied it to all our experiments in infsection
80.0
Test Accuracy
Test Accuracy
77.5
60
75.0
/2.5
70.0
glabal decay wio global docay
w/b local Ir decay
67.5
10-3
10-2
0.05
0.1
0.5
Initial learning ratc
a
Figure 6: FEDAVG with ConvNet Step-non- i.dd CIFAR- 10 with without learning rate decay within each round of local training.
Figure 7: FEDAVG with ConvNet on Dirichlet- non-i.i.d CIFAR- 10 with or without learning rate decay latter rounds of communication. We ex- perimented with different values a in Dir(a).
B.3 EFFECTS OF WEIGHT DECAY IN LOCAL CLIENT TRAINING
Federated learning on non-i.i.d data clients prone to model drift due to the deviation of local data distributions and is sensitive to the number of local epochs E (Wang al. 2020 Li et al. 2020b McMahan et al.) 2017). To prevent local training from over-fitting the local distribution, we apply lg regularization weight decay. In different context of distributed deep learning, Sagawa et al. (2020) also showed that E, regularization can improve the generalization by preventing the local model from perfectly fitting the non-i.i.d training data.
15

---

Published as a conference paper at ICLR 2021
Table 7: FEDBE with models sampled from Dirichlet distribution on Step-non-ii.d. CIFAR-IO. We compare different a x I in setting the parameter of a Dirichlet distribution.
ConvNet
ResNet20
0.1
72.5+0.44
75.9=+0.66
0.5
73.6+0..3
0..3+0.86
1
74.2+:051
77.1=10.11
2
73,240.83
76.8=10.55
As shown in Figure 8 where we compared FEDAVG and FEDBE with without weight decay in local training, we found that weight decay not only leads to higher test accuracy but also makes both algorithms more robust to the choice local epochs E.
74
Test Accuracy
72
70
68
FedBE w/ weight decay
FedBE wlo weight decay
FedAvg w/ weight decay
66
FodAvg wo weight decay
10
20
50
100 # Local epochs
150
Figure 8: FEDAVG and FEDBE with ConvNet on CIFAR- 10 (Step-non-i...d) with or without weight decay in local training, for different numbers of local epochs E.
C EXPERIMENTAL RESULTS AND ANALYSIS
C.I GLOBAL MODEL SAMPLING IN FEDBE
In the main paper, mainly report accuracy by FEDBE with models sampled from Gaussian distribution (cf. [qquation of the main paper). Here we report results using a Dirichlet distribution for model sampling (cf. Equation of the main paper) in Table 7 We compare different in setting the parameter change of Compared to Table I of the main paper, FEDBE with Dirichlet is slightly worse than FEDBE with Gaussian (by Σ 0.5%) but much better than FEDAVG.
x 1 Dirichlet distribution. We see that the accuracy is not sensitive to the
To further study the models sampled from the global model distribution constructed in FEDBE, compare the prediction accuracy and confidences (the maximum values of the predicted probabil- ities) of clients' models and sampled models. We show the histogram of correctly and incorrectly predicted test examples at different prediction confidences. As shown in Figure 9 we observe that clients' models tend to be over-confident by assigning high confidences wrong predictions. We hypothesize that it is because clients' local training data are scarce and class-imbalanced. On the other hand, sampled models have much better alignment between confidences and accuracy (i.e., higher confidences, higher accuracy).
C.2 ANALYSIS WEIGIT AVERAGE, (BAYESIAN) MODEL ENSEMBLE, AND DISTILLATION
To investigate the difference of weight average, (Bayesian) model ensemble, and distillation in making predictions, we focus on one-round federated learning. in which the client models are trained the same way regardless of what aggregation approach to be used. We experiment with Step-non-i.i.d. CIFAR- 10 using ConvNet, and train the 10 local client models for 200 epochs. We then compare (a) weight average combine the models, (b) model ensemble, and (c) Bayesian model ensemble
16

---

Published as a conference paper at ICLR 2021
Swiss Rolls: Clients
Swiss Rolls: Samples
140
Correct Incorrect
120
100
80
60
60
40
0.0
0.2
0.4
0.6
0.0
1.0
0.2
0.4
0.6
1.0
CIFAR-10: Clients
CIFAR-10: Samples
2300
2500
2000
1500
1000
300
0.2
0.5
13.6
0.5
0.6
0.7
0.8
0.5
1.0
0.4 0.6 0.5
Figure Histograms correctly and incorrectly predicted examples (vertical axes) along the confidence values (the maximum values of the predicted prohabilities). Upper row: Swiss roll dataset used in Figure of main paper (averaged over 3 clients or sampled models); lower row: Step-non-ii.d. CIFAR- (averaged over 10 clients sampled models).
Table 8: One-round federated learning on Step-non-i.i.d CIFAR-10 with ConvNet. We compare different strategies to combine the clients' local models, including weight average, (Bayesian) model ensemble, and ensemble distillation (with SGD or SWA)
Method
Distillation
None
SGD
SWA
(a) Weight average
2..7
(b) Model ensemble
60.5+1028
32.0
33.1+1.02
(c) Bayesian model ensemble
62.5++0.35
35.16
35.7++0.86
with M 10 extra samples beyond weight average and individual clients. For (b) and (c), we further apply knowledge distillation using the unlabeled server data to summarize the ensemble predictions into a single global model using SGD or S SWA. We note that, method (a) is equivalent to one-round FEDAVG; method (b) without distillation is the same as I-Ensemble; method (b) with SGD distillation is equivalent to one-round y-Distillation; method (c) with SWA is equivalent to one-round F.DBE. Table 8 shows the results with several interesting findings. First, without distillation, model ensemble clearly outperforms weight average and Bayesian model ensemble further adds a 2% gain, supporting our claims in |section 3 Second, summarizing model ensemble into one global model largely degrades the accuracy. showing the challenges of applying ensemble distillation in federated learning. The 60.5% and 62.5% accuracy by ensemble, although relatively higher than others, may still be far from perfect to be used as distillation targets. Third, distillation with SWA outperforms SGD for both model ensemble and Bayesian model ensemble, justifying our proposed usage of SWA. Fourth, with without distillation. Bayesian model ensemble always outperforms model ensemble, with very little cost of estimating the global model distribution and performing sampling. Fifth, with the degraded accuracy after distillation, model ensemble and Bayesian model ensemble. after distilled into a single model, still outperforms weight average notably. Finally, although in the one-round setting hardly the advantage of distilling the model ensemble into single model,
17

---

Published as a conference paper at ICLR 2021
FEDAVG: inference CIFAR-10
FEDAVG: inference CIFAR-100
FEDBE: inference on CIFAR-10
FEDBE: inference on CIFAR-100
Figure 10: Feature visualization of FEDAVG and FEDBE models. All models are trained on Step-non i. CIFAR-10, 10. then inference on CIFAR-10/100 test sets. The features are colored with the ground-truth labels (CIFAR-10) 10) the predictions (CIFAR-100).
with multiple rounds of communication as in the main paper (cf. (subsection 5.2), its advantage becomes much clear-it allows the next-round local training to start from better initialization (in comparison weight average) and eventually leads to much higher accuracy than I-Ensemble.
C.3 FEDBE vs. FEDAVG
We provide further comparisons between FEDBE and FEDAVG, First, we perform Bayesian model ensemble at the end of FEDAVG training (Step-non-ii.d. CIFAR-10), We achieve 72.5% with ResNet20, better than FEDAVG (70.2% in Table of the main paper) but still worse than EEDBE (77.5%), demonstrating the importance of incorporating Bayesian model ensemble into multi-round federated learning.
To further analyze why FEDBE improves over FEDAVG, we train models Step-non-i.i.d CIFAR-10. inference on the test sets, and visualize the features. We trained a FEDBE ResNer20 model for only 15 rounds such that the test accuracy is similar a FEDAVG ResNet20 model trained for 40 rounds. In Figure 10, we plot their features using t-SNE (Maaten & Hinton 2008) on the CIFAR- 10/100 testing sets (consider 3 semantically different classes: automobile, cat, and frog) and color the features with the ground-trunt -ruth labels of CIFAR- 10 test set or the predictions of the models on CIFAR- 100 test set. Interestingly. we observe that the features FEDBE are more discriminative (separated) than the features of FEDAVG, especially on CIFAR-100, even if FEDBE and FEDAVG have similar test accuracy on CIFAR-10.
We further discuss Table 3 of the main paper. We perform data augmentation on T in knowledge distillation. This explains why we obtain improvement over FEDAVG when using the FEDAVG predictions as the target labels (72.6% 77..%%%/33%% vs. 72.0% 70..2% in Table of the main text, using ConvNet/ResNer20), We note that without data augmentation, using FEDAVG predictions as the target leads to gradients in knowledge distillation since initialize the student model with FEDAVG. The results suggest the slight benefit of collecting unlabeled data for knowledge distillation in model aggregation.
18

---

Published as a conference paper at ICLR 2021
80
Test Accuracy
78
76
74
72
FedBE-SWA
FedHH--SGl
70
10
30
# #Epochs
Figure 11: ResNet20 test accuracy on Step-non-i.i.d. CIFAR-I0, with different numbers of epochs for distillation using SGD SWA for FEDBE.
C.4 FEDBE WITH SWA AND SGD
We found that distillation with SGD is more sensitive to noisy labels and the number of epochs. For ResNet20 (Talle 1). FEDBE with FEDAVG C w/o SWA (i.e., using SGD) achieves 74.9% with 20 epochs but 74.0% with 40 epoch. In contrast, FEDBE with SWA is much stable. As shown in Figure 11 the accuracy stays stable with more than 10 epochs being used, achieving 77.5% with 20 epochs and 77.3% with 40 epochs.
C.5 BATCH NORMALIZATION GROUP NORMALIZATION
Hsieh et (020)) showed that FEDAVG non-i.i.d cases can be improved by replacing the batch normalization (BN) layers with group normalization (GN) layers (Wu & Hej 2018). However, we observe that ResNets with GN converge much slower, which is consistent with the observations in (Zhang et al. 2020). In our CIFAR-10 (Step) experiments, FEDAVG using ResNet20 with GN can outperform that with BN slightly both are trained with 200 rounds (76.4% vs. 74.6%). FEDBE can further improve the performance: FEDBE with GN/BN achieves (9.6%/80.22
C.6 COMPATIBILITY WITH SCAFFOLD
SCAFFOLD (Karimireddy al. 2020 is recently proposed FL method to regularize local training. We experiment with SCAFFOLD on non-i.i.d. (Step) CIFAR-10. We find that SCAFFOLD cannot directly perform well with deeper networks (ResNer20: 59.4%; ResNet32: 55.3%). Nevertheless, FEDBE +SCAFFOLD can improve upon it, achieving 76.4% and 72.7% respectively.
To further analyze why FEDBE improves SCAFFOLD, we plot the test accuracy SCAFFOLD vs. FEDBE +SCAFFOLD at every communication round. We see that both methods perform similarly in the early rounds. SCAFFOLD with weight average could not improve the accuracy after roughly 10 rounds. FEDBE +SCAFFOLD, in contrast, performs Bayesian ensemble and distillation to obtain the global model, bypassing weight average and gradually improving the test accuracy. We therefore argue that, long FEDBE can improve SCAFFOLD slightly at every later round, the ultimate gain can be large. We also attribute the gain brought by FEDBE to the robustness of Bayesian ensemble for model aggregation.
C.7 PARTIAL PARTICIPATION ON CIFAR- 100 (CF. SUBSECTION 5.3)
We also conduct the experiments on CIFAR-100 with the non-i.i.d. Step setting. We consider setting with 100 clients, in which 10 clients are randomly sampled at each round and iterates for 100 rounds, similar to MMMMahan et al. 2017). We split 10K images from the 50K training images to the server, and distribute the remaining ones the clients. Each client has 5 major classes (61 images each) and 95 minor classes (1 image each). Table shows the results: FEDBE consistently outperforms FEDAVG.
19

---

Published as a conference paper at ICLR 2021
80
70
Test Acc
60
50
40
30
SCAFFOLD
SCAFF-D
20
10
20
30
40
Rounds
Figure 12: Step-non-i.idd CIFAR-10 experiments accuracy curves of SCAFFOLD on ResNet20.
Table 9: Non-i.id CIFAR-100
Method
ConvNet
ResNet20
ResNer32
FEDAVG FEDBE
32.5+-078
37.5++0.55
33.3++055
36.6:10.22
43.5:10.99
37.7110.69
D DISCUSSION
D.1 EXTRA COMPUTATION COST
computation compared FEDAVG. The extra cost is the server and no extra burden is on the clients. In practice, the server assumed to be computationally rich so the extra training time is negligible w.r.t. communication time. Using a 2080 Ti GPU CIFAR- 10 (ConvNet), building distributions and sampling takes 0.2s, inference of a model takes 2.4s, and distillation takes 10.4s. Constructing the ensemble predictions T
FEDBE involves
(()) where P MT E 2== p(ylar;; (T))), requires each (mm) to be evaluated U, which can be easily parallelized in modern GPU machines. The convergence speed of the Monte Carlo approximation in Equation 4 1/VM, yet we observe that A 10 ~ 20 is sufficient for Bayesian model ensemble to be effective.
D.2 FEDAVG ON DEEPER NETWORKS
Deeper models are known be poorly calibrated (Guo et al. 2017), especially when trained on limited and imbalanced data. The loss surfaces can non-convex (Garipov et al.] [2088] Draxler et al. 2018). FEDAVG thus may not fuse clients well and may need significantly more rounds of communication with local training of small step sizes to prevent client's 's model drifting.
E FURTHER ANALYSIS
E TEST ACCURACY AT DIFFERENT ROUNDS
We follow the experimental setup in section 5 and further show the test accuracy of the compared methods at different communication rounds (in total 40 rounds) in Figure 13 Specifically, we experiment with ResNet20 and ResNet32 for both the Step-non-i.i.d. and Dirichlet-non-iiid. settings on CIFAR- 10 (the results at 40 rounds are the same as those listed Table I.. FEDBE obtains the highest accuracy after roughly 10 rounds, and could gradually improve more rounds are involved. Interestingly, v-Distillation, which performs ensemble directly over clients' models without other sampled models, normally obtains the highest accuracy in the first 10 rounds, but is surpassed by FEDBE after that. We hypothesize that the first 10 rounds, as the clients models are still not well trained, the constructed distributions may not be stable. We also note that except 1-Ensemble,
20

---

Published as a conference paper at ICLR 2021
Step-non-ii.d, ResNet20
Dirichlet-non-iiid, ResNet20
80
80
70
Ensemble
Test Acc
Redlyg
Test Acc
60
DDstillatonn
60
ADistilation
SWA)
FodBE two SWA
reddtt
50
40
40
0
10
20 Rounds
30
40
Rounds
Step-non-ii.dd ResNet32
Dirichlet-nonni.i.d ResNet32
80
70
Test Acc
60
Rad4ig
Ensembie
60-
TecAug
DDistilation
*Distilatior
Test
50
SWA)
50
reditt
FecBE
40
40
30
0
10
20
30
40
10
20
30
40
Rounds
Rounds
Figure 13: CIFAR-IO curves of test accuracy at different communication rounds. We study both non-i.i.d
settings (Step and Dirichlet) using ResNet20 and ResNet32 (cf. subsection E.
the other methods with ensemble mostly outperform FEDAVG the first 10 rounds, showing their robustness in aggregation.
21


Reviews:
[{"review_id": "s3NX3XnWhqc", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "Summary:\n\nThe paper proposes a new model aggregation approach named FedBE in federated learning. In each round, FedBE first fits the local models into a Gaussian/Dirichlet distribution as a posterior distribution of the global model. Then, FedBE uses the sampled models from the distribution, the local models, and the averaged model as an ensemble to label a public dataset. The labeled dataset is used to train a global model. The experiments show that FedBE can outperform FedAVG on CIFAR-10 using ConvNet and ResNet.\n\nPros:\n\n(1) The paper is well written and easy to follow.\n\n(2) The studied problem is important and the proposed idea is interesting.\n\n(3) The experiments are comprehensive and cover many aspects.\n\nI have the following concerns, mainly about the experiments.\n\n(1) In Algorithm 1, only the sampled models are used to predict the unlabeled dataset. However, in the experiments, the sampled models, the local models, and the averaged model are all used as an ensemble. Please keep consistency. Also, if the local models may perform badly due to client’s model drift, why the local models are also included in the ensemble?\n\n(2) In Section D.6, SCAFFOLD has a bad test accuracy. I don’t understand why FedBE can improve SCAFFOLD a lot. If the averaged model is bad in SCAFFOLD, the sampled models are also very likely to be bad. However, the authors show that the accuracy of FedBE+SCAFFOLD is still close to FedBE+FedAvg. Can the authors explain it?\n\n(3) From Table 7, FedBE with Dirichlet distribution is worse than FedBE with Gaussian distribution although it introduces an additional parameter $\\alpha$. Is there any explanation?\n\n(4) The size of unlabeled public data is large in the experiments. The paper uses CIFAR100+tiny-imagenet as the public dataset, which has a much larger size than the training dataset (i.e., CIFAR10). However, in reality, the data are usually private and the server may not be able to collect a large related public dataset. Can the authors show the performance of FedBE with a smaller public dataset?\n\n(5) The authors can add a figure to show the test accuracy of different approaches versus the number of communication rounds. It is more clear to see the performance of different approaches with such a figure.\n\n(6) The paper presents a simple averaging version in the main text and a weighted averaging version in Appendix. Which version is used in the experiments?\n\nMinor issue: Figure 1 uses Dirichlet distribution as an example. Since the paper focuses on the Gaussian distribution, the authors can present an example based on Gaussian distribution instead.\n\n\n==========Post rebuttal============\n\nThanks for the authors' response. The authors have addressed some of my concerns. I have raised my score accordingly.", "extra_scores": {}, "extra_text": {}}, {"review_id": "lxirOwo3opu", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "-- Summary --\nThe paper proposes a new approach to model aggregation at each round performed on the server using a combination of Bayesian ensembling, model distillation, and stochastic weight averaging. The proposed method is essentially a drop-in replacement for weighted averaging based model aggregation used by FedAvg. The new method, FedBE, is extensively evaluated experimentally on federated image classification datasets where the authors show improved performance over multiple baselines.\n\n\n-- Overall evaluation --\nI find the paper well written and the idea of using ensembles and distillation for model aggregation in FL settings very interesting. Although the overall idea is identical to FedDF by Lin et al. (2020) who proposed to use ensembling and distillation and the authors cite and compare with that work, I think forming the ensemble distribution, sampling from it, as well as using SWA for distillation are all novel components; the contribution of each introduced component is empirically analyzed in ablation studies. Overall, while I don't find the approach extremely novel or insightful, it seems like a nice contribution to the literature. Below are some comments and questions that I hope will help further clarify and improve the paper.\n\n\n-- Comments and questions --\n\n- Section 3.3: While I see how different distributions over w are formed (Gaussian or Dirichlet), I'm not sure I understand why these distributions approximate the actual posterior? In other words, is it true that the distributions described by the authors minimize KL (or some other) divergence to the true posterior? Can you justify that?\n\n- Regarding the Dirichlet (eq. 7), how are the gamma parameters estimated? Or are these simply hyperparameters?\n\n- It's nice to have a toy example in section 3 (Figure 1) that illustrates the approach. However, some details are missing from the description. For example, how alphas were selected? How many clients per round were used? How did the test data distribution differ from the data on the clients?\n\n- Computational complexity: I wonder how much is the computational overhead one has to pay to run FedBE on the server instead of FedAvg? In some practical settings, FedBE could be a great replacement for FedAvg, but in other settings, it might be too compute-intensive. It would be nice if the authors can discuss tradeoffs and analyze computational complexity.\n\n- Related to the previous question: I wonder if distillation at each round is actually necessary? Given an ensemble (which is relatively cheap to obtain), why not sample multiple w's from it and initialize local models on clients using these samples? Do the authors expect it to work worse?\n\n\n-- Experiments --\n\n- Table 1 clearly demonstrates that FedBE dominates other methods on CIFAR-10. However, I'm puzzled by the reverse trends for centralized SGD and FL with the increase in the complexity of the model: more complex models trained with FL have worse performance than simpler models, and the opposite is true for centralized training. I wonder if ResNets were actually undertrained to some extent? Can the authors explain this phenomenon?\n\n- Section 5.1, paragraph 2: how come the number of local epochs is computed as E = K / batch_size? Also, which metric was used exactly to tune the number of local epochs for FedAvg? After how many rounds the performance of FedAvg was evaluated and used for selecting E?\n\n\n-- Minor --\n\n- Is there any empirical evidence that supports the claim in footnote 2?\n\n- Eq. 2 assumes that all clients have local datasets of the same size. Appendix A (eq. 10) provides the correct expression, but why not give that expression directly in the main text? The difference is literally 2 symbols. Same for eqs. 11-13.", "extra_scores": {}, "extra_text": {}}, {"review_id": "b4iWPELZWO", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "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 focuses on improving model aggregation on the server of multi-round federated learning. For this purpose, this paper proposes FedBE, which combines the idea of Bayesian model ensembling and distillation based on stochastic weight averaging. A set of unlabelled data on the server is required for distillation. Experiments are conducted on CIFAR-10/100 and tiny-ImageNet with both IID and non-IID cases.\n\nThe paper is well written and easy to follow. Though Bayesian model ensemble and distillation based on SWA are not new topics, using the idea of them in the field of federated learning is novel to the best of my knowledge. \n\nMy major concern is about the experiment setting for using unlabeled data. According to my understanding, from the main body of the paper (Sec 5.2), 10K training images are used as unlabeled data saved on the server and the rest is distributed to each client. For baseline method like FedAvg, these 10K images are useless since it does not need unlabeled data for model average on the server. The different number of total training instances between FedBE (40K labeled + 10K unlabeled) and FedAvg (40K labeled) makes the comparison unfair. \n\nBesides, the analysis about either computational complexity or empirical wall clock time is needed. FedBE replaces the original model average with ensemble prediction & training, which is very slow. And these computations on the server seem cannot be parallelized with the computation on the client.  As a result, the computation on the server maybe the bottleneck of FedBE. So providing detailed analysis is important for better evaluating FedBE.", "extra_scores": {}, "extra_text": {}}, {"review_id": "VF0ulVSowgO", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "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": "The paper presents FEDBE, a novel method for federated learning. Models are trained locally in the clients' side then communicated to the server. The global model is constructed by Bayesian model ensemble. Knowledge distillation is used to construct the model that is passed to the clients for the next round. FEDBE is simple and shows strong empirical performance.\n\nStrong points:\n\n- The method demonstrates strong performance versus common baselines\n- Extensive experiments were conducted.\n- The paper is well written.\n\n\nClearly state your recommendation (accept or reject) with one or two key reasons for this choice.\n- I recommend to accept this paper.\n- This type of work is definitely needed to enable more powerful models that preserves the privacy of the clients.\n\n\nQuestions:\n\n- What is the effect of number of epochs between model collections (25 steps vs. 50 or more)? Is this what is shown in Figure 10? If so please clarify more.\n- What is the effect of the step size decaying on SWA? IS the method sensitive to this choice?\n- What are the hyperparameters that the method is sensitive for?", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This is a well written paper with good experimentation.  The paper builds on the work of FedDF and does ablation studies to demonstrate its improvements.  The key original idea is the use of a common pool of unlabeled data which is used in transmitting partial results between local and global servers.  The results seem pretty good.

From a practical viewpoint, the unlabelled common data will, in most cases, need to be generated/artificial data since it will need to be public (to the other servers at least).  This option should be tested to demonstrate feasability.

AnonReviewer2 was concerned about whether it was fair to provide additional unlabelled data.  The authors tested this out and showed it was OK.  Regardless, the different servers could easily generate artificial data for this purpose.  AnonReviewer1 had a number of issues which the authors largely addressed. The other two reviewers appreciated the paper.  All reviewers gave constructive suggestions.

Decision: Accept (Poster)