A Hierarchical Bayesian Approach to Federated Learning
Authors: Minyoung Kim
Timothy Hospedales
Venue: Submitted to ICLR 2023 (2023)
Keywords: Federated Learning
Bayesian Methods
Probabilistic Models

Abstract: We propose a novel hierarchical Bayesian approach to Federated learning (FL), where our models reasonably describe the generative process of clients' local data via hierarchical Bayesian modeling: constituting random variables of local models for clients that are governed by a higher-level global variate. Interestingly, the variational inference in our Bayesian model leads to an optimization problem whose block-coordinate descent  solution becomes a distributed algorithm that is separable over clients and allows them not to reveal their own private data at all, thus fully compatible with FL. We also highlight that our block-coordinate algorithm has particular forms that subsume the well-known FL algorithms including Fed-Avg and Fed-Prox as special cases. That is, we not only justify the previous Fed-Avg and Fed-Prox algorithms whose learning protocols look intuitive but theoretically less underpinned, but also generalise them even further via principled Bayesian approaches. Beyond introducing novel modeling and derivations, we also offer convergence analysis showing that our block-coordinate FL algorithm converges to an (local) optimum of the objective at the rate of $O(1/\sqrt{t})$, the same rate as regular (centralised) SGD, as well as the generalisation error analysis where we prove that the test error of our model on unseen data is guaranteed to vanish as we increase the training data size, thus asymptotically optimal.

Under review as a conference paper at ICLR 2023
A HIERARCHICAL BAYESIAN APPROACH TO FEDER- ATED LEARNING
Anonymous authors Paper under double-bling review
ABSTRACT
We propose a novel hierarchical Bayesian approach to Federated Learning (FL), where our model reasonably describes the generative process of clients' local data via hierarchical Bayesian modeling: constituting random variables of local models for clients that are governed by higher-level global variate. Interestingly, the variational inference in our Bayesian model leads to an optimisation problem whose block-coordinate descent solution becomes distributed algorithm that is separable over clients and allows them not to reveal their own private data at all, thus fully compatible with FL. We also highlight that our block-coordinate algorithm has particular forms that subsume the well-known FL algorithms including Fed-Avg and
as special cases. That is, we not only justify the previous Fed-Avg and Fed-Prox algorithms whose learning protocols look intuitive but theoretically less underpinned, but also generalise them even further via principled Bayesian approaches. Beyond introducing novel modeling and derivations, we also offer convergence analysis showing that our block coorddcoo FL algorithm converges to an (local) optimum of the objective at the rate of (1/vVt, the same rate as regular (centralised) SGD, as well as the generalisation error analysis where prove that the test error of our model on unseen data is guaranteed to vanish as we increase the training data size, thus asymptotically optimal.
INTRODUCTION
Federated Learning (FL) aims to enable set of clients to collaboratively train a model in privacy preserving manner, without sharing data with each other central server. Compared conventional centralised optimisation problems, FL comes with a host of statistical and systems challenges such as communication bottlenecks and sporadic participation. The key statistical challenge is non-i.i.d. data distributions across clients, each of which has different data collection bias and potentially a different data annotation policy/labeling function for example, in the case of any user preference learning. The classic and most popularly deployed FL algorithms are Fed-Avg (McMahan et al.) 2017) and Fed-Prox (Li et al. 2018), however, even when global model can be learned, it often underperforms on each client local data distribution in scenarios of high heterogeneity (Li et al., 2019: Karimireddy et al., 2019 Wang et al. 2020). Studies have attempted to alleviate this by personalising learning at each client, allowing each local model to deviate from the shared global model (Sun et al. (2021). However, this remains challenging given that each client may have a limited amount of local data for personalised learning.
These challenges have motivated several attempts to model the FL problem from a Bayesian perspec- tive. Introducing distributions model parameters 0 has enabled various schemes for estimating global model posterior p(0) Din)) from clients' local posteriors p(0) D,), or to regularise the learning of local models given a prior defined by the global model Zhang et al. (2022); Al-Shedivat al. (202)) Chen & Chao 202T). However, these methods are not complete and principled solutions having not yet have provided full Bayesian descriptions of the FL problem, and having had resort to ad-hoc treatments to achieve tractable learning. The key difference is that they fundamentally treat network weights 0 as a random variable shared across all clients. We introduce a hierarchical Bayesian model that assigns each client it's own random variable for model weights 0,, and these linked via higher level random variable ó p(f1:N, o) crucial benefits: Firstly, given this hierarchy, variational inference in our framework decomposes into separable optimisation problems over 0,s and o, enabling practical Bayesian learning algorithm
IT p(0:lo). This has several

---

Under review as a conference paper at ICLR 2023
0
(a) Overall model
(b) Individual client (c) Global prediction (d) Personalisation
Figure I: Graphical models. (a) Plate view of iid clients. (b) Individual client data with input images given and only p(y(æ) modeled. (c) & (d): Global prediction and personalisation as probabilistic inference problems (shaded nodes evidences, red colored nodes targets infer, test input in global prediction, DP training data for personalisation and x? test input).
to be derived that is fully compatible with FL. constraints, without resorting ad-hoc treatments or strong assumptions. Secondly, this framework can be instantiated with different assumptions on p(0:ló) to deal elegantly and robustly with different kinds of statistical heterogeneity, well for principled and effective model personalisation.
Our resulting algorithm, termed Federated Hierarchical Bayes (FedHBB is empirically effective, as we demonstrate in wide range of experiments on established benchmarks. More importantly, it benefits from rigorous theoretical support. In particular, we provide convergence guarantees showing that FedHB has the same (l/vVT) convergence rate as centralised SGD algorithms, which are not provided by related prior art /Zhang et al. (2022): [Chen & Chao (202T). We also provide generalisation bound showing that FedHB is asymptotically optimal, which has not been shown by prior work such as Al-SShedivat et al. (2021). Furthermore we show that FedHB subsumes classic methods FedAvg McMahan al. (2017 and FedProx Li et al. (2018) as special cases, and ultimately provides additional justification and explanation for these seminal methods.
BAYESIAN FL: GENERAL FRAMEWORK
We introduce two types of latent random variables, e and {0.). Each 0; is deployed the network weights for client i's backbone. The variable & can be viewed as globally shared variable that is responsible for linking the individual client parameters 0,. We assume conditionally independent and identical priors, p(0 (o) II p(0.lo). Thus the prior for the latent variables (ø, (0, ) is formed hierarchical manner as 11. The local data for client denoted by D,, is generated by O.,
(Prior) p(6, (I.) p II p(0
(Likelihood) p(D.10,)
II plyla, 0,),
(1)
[x.,)ED,
where p(y (x, is conventional neural network model (e.g., softmax link for classification tasks). See the graphical model in Fig. (a) where the iid clients are governed by a single random variable ø.
Given the data DI..... DN we infer the posterior, p(d, 01.N DI.N) p(o) II p(0, ())(( which is intractable general, and we adopt the variational inference to approximate it:
g(o, :I.N; = q(d; Lo) II (0;; Le),
where the variational parameters consists of La (parameters for (())) and {L:) 's (parameters for qi (0:)'s from individual clients). Note that although 0,'s independent across clients under (2). they are differently modeled (emphasised by the subscript in notation 9:). reflecting different posterior beliefs originating from heterogeneity of local data D.'s.
2.1 FROM VARIATIONAL INFERENCE TO FEDERATED LEARNING ALGORITHM
Using the standard variational inference techniques Blei et al. 2017 Kingma & Welling 2014 can derive the ELBO objective function (details in Appendix A). We denote the negative ELBO by L
Note that we do not deal with generative modeling of input images v. Inputs x always given, and only conditionals p(y)a)) modeled. See Fig. b) for the in-depth graphical model diagram.
2

---

Under review as a conference paper at ICLR 2023
(to be minimised over L) as follows:
((L) Σ (Eu.coll (oo.oo)-logpD..10.) Bggo))))))))))))))))))) (Kt))][K(( [KL(((()))) + (L(a()))l(()) (3)
where we drop the dependency on L in notation for simplicity. Instead of optimizing (3) over the parameters L jointly as usual practice, we consider block-wise optimisation, also known block- coordinate optimisation (Wright 2015), specifically alternating two steps: (i) updating/optimizing all
L E.'si: 1...., N while fixing Lo. and (ii) updating Lo with all Lç's fixed. That is,
Optimisation over Li...., LN (Lo fixed).
N
min {Le)-
-o log D. + (() (L)(((()))l((0]0)))
(4)
As (4) is completely separable over i, and we can optimise each summand independently as:
min Ci(L,) (0,40 log p( [.)] + + ( [g, La)l:l [))lp(([])))] (5)
So
( constitutes local update/optimisation for client i. Note that each client i needs to access private data D. only without data from others, thus fully compatible with FL.
Optimisation over Lo (L1...., LN fixed).
N Ealorto Ee
min Lo(Lo) == Lo)lp(()) Lo
(6)
1=1
This constitutes server update criteria while the latest (0;; L.)'s from local clients being
fixed. Remarkably, the server needs not access any local data all, suitable for FL. This
nice property originates from the independence assumption in our approximate posterior ( Interpretation. First, server's loss function (6) tells us that the server needs to update q(:; Lo) in such a way that (i) it puts mass on those that have high compatibility scores log p(0,(00
(o) with the
current local models 0g ~ qi (0.), thus aiming to be aligned with local models, and (ii) it does not deviate from the prior p(ő). Clients' loss function (5) indicates that each client needs minimise the class prediction error on its own data D, (first term), and at the same time, to stay close to the current global standard @ ~ g( by reducing the KL divergence from p(0, ()) (second term).
2.2
FORMALISATION OF GLOBAL PREDICTION AND PERSONALISATION TASKS
Two important tasks in FL are: global prediction and personalisation. The former evaluates the trained model on novel test data sampled from a distribution possibly different from training data. Personalisation is the task of adapting the trained model on a new dataset called personalised data. In Bayesian model, these tasks can be formally defined Bayesian inference problems.
Global prediction. The task is to predict the class label of a novel test input x² which may or may not come from the same distributions the training data Dı,... ... D Under can be turned into probabilistic inference problem p(y" |xx, DI.N). Let 0 be the local model that generates the output y² given x². Exploiting conditional independence from Fig. [c),
Bayesian model, it
1] pl"" p(0 p(o) d0ds
plu |z" D1.N)
2 =If plo"la" (000')))) p(0 q()) d0dø = la" 0)
(DDOTO) g(o)do d0,
(8)
where (8) we use p()) DEN) ss (()) The inner integral (in parentheses) in (8) either admits a closed form (Sec. 3.1)
be approximated (e.g.. Monte-Carlo estimation).
Personalisation. It formally refers to the task of learning a prediction model p(y|z) given an unseen (personal) training dataset DP that comes from some unknown distribution pd p² y). that the personalised model ý performs well on novel (in-distribution) test points (xx y²) ~ (x, y). Evidently need to exploit (and benefit from) the trained model from the FL training stage. To this end many existing approaches simply resort to finetuning, that is, training on DP warm-starting with the FL-trained model. However, potential issue is the lack of solid principle how to balance the initial FL-trained model and personal data fitting to avoid underfitting and overfitting. In our Bayesian framework, the personalisation can be seen as another posterior inference problem with additional evidence of the personal training data DP, Prediction on a test point x" amounts to inferring:
ply" DP, DI.N) f p(y" x², .)) p00 D1:N) dA.
(9)
3

---

Under review as a conference paper at ICLR 2023
So, it boils down to the task of posterior inference p(0 D1.N) ) given both the personal data DP and the FL training data DI.N- Under our hierarchical model, by exploiting conditional independence from graphical model (Fig. (()),, we can link the posterior to our FL-trained (()) as follows:
p(0) D", o) p(o DEN) do x f ( D", ( do R p(0 (10)
D(0 DI.N) s
where we disregard the impact of D2 on the higher-level given the joint evidence, p(ø) DP, DI:N) R Pi.) due to the dominance of D1.N compared to smaller DP. The last part of (TO) makes approximation using the mode o* of q()). which is reasonable for our two modeling choices for
(()) to be discussed in Sec. 3.1 and Sec. 3.2 Since dealing with DP, Φ²) involves difficult marginalisation p( (D) = pl ()) Φ )do, we adopt variational inference, introducing a tractable variational distribution u(f) p(0)DP, x²). Following the usual variational inference derivations, we have the negative ELBO objective (for personalisation):
min Ex log p(D + (L((v()]Ip(())) (l)())))))
(11)
Once we have the optimised model v, our predictive distribution becomes:
S
12" DP, (I,N)
Σ ply 12, g(t)), where g(*) ~ v(0).
(12)
which simply requires feed forwardigg test input x through the sampled networks ((5) and averaging.
Thus far, we have discussed general framework, deriving how the variational inference for our Bayesian model fits gracefully in the FL problem. In the next section, we define specific density
families for the prior (p()), p(O)lo)) and posterior (g(Ø), (B.)) as our proposed concrete models.
BAYESIAN FL: Two CONCRETE MODELS
We propose two different model choices that we find the most interesting: Sariaal-nvere-WWiii (Sec. 3.1 and Mixture (Sec. 3.2) To avoid distraction, we make this section concise putting only the final results and discussions, and leaving all mathematical details in Appendix B and C
3.1 NORMAL- WISHART (NIW) MODEL
We define the prior as a conjugate form of Gaussian and Heemal-lnverse-ishheee With d = (1,, E), p(0)= = NIW E; = NN; po, A)) IW(E; 20),
(13)
p(0 o) (0 E), i = .N, where A = {Ho, Co Ao } is the parameters of the NIW. Although A can be Iearned via data marginal likelihood maximisation (e.g., empirical Bayes), but for simplicity we leave it fixed a No = 0, Σg I, Ao = 1, and Pn = d + 2 where d is the number of parameters in 0, μ. Next, our choice of the variational density family for q()) is the NIW, not just because it the most popular parametric family for pair of mean vector and covariance matrix 6 (J,, E). but it can also admit closed-form expressions in the ELBO function due to the conjugacy as we derive in Sec. B.I
(14)
q(0)= NIW (mo, Vo, ln m0)) N N(xx; ma, 1 IWV Vo, no)
(15)
Although the scalar parameters la. no can be optimised together with ma. Vo. their impact is less influential and we find that they make the ELBO optimisation a little bit cumbersome. So we fix lo no with some near-optimal values by exploiting the conjugacy of the NIW under Gaussian likelihood (details in Appendix BB., and regard mg, Vn as variational parameters, Lo = {mo, VoJ. We restrict Vo to be diagonal for computational tractability. The density family for q: (0.)'s can be Gaussian, but we find that it is computationally more attractive and numerically more stable to adopt the mixture of two spiky Gaussians that leads to the MC-Dropout (Gal & Ghahramani] 2016). That is,
qi (0.) II (p- [[(,[ m. m,111, (2) + (1 p) (00,(0)0,,
(16)
where (i) my is the only variational parameters (L. f). (ii) [[] indicates a column/layer in neural network parameters where goes over layers and columns of weight matrices, (iii) p is the (user-specified) hyperparameter where 1 - p corresponds to the dropout probability, and (iv) e is small constant (e.g., 10-4) that makes two Gaussians spiky, close to the delta functions.
Client update. We apply the general client update optimisation 5) to the NIW model. Following the approximation of (Gal & Ghahramani] 2016) for the KL divergence between a mixture of Gaussians
"This choice ensures that the mean of E Eequal T, and j is distributed as O-mean Gaussian with covariance E.
4

---

Under review as a conference paper at ICLR 2023
166 and Gaussian 14). we have the client local optimisation (details in Appendix B
min Li (m,.) := logp(D;(mm.) + (no + d + 1)(m: - -mo) Vo- (m; - mo),
(17)
2
where m, is the dropout version of m,, i.e., a reparametrised sample from [6). Note that my and Vo are fixed during the optimisation. Interestingly (17) generalises Fed-Avg (McMahan et al. (017)) and
Fed-Prox (Li et al. 2018): With p (i.e.. dropout) and setting Vi al, [17 reduces to the client update formula for Fed-Prox where constant a controls the impact of the proximal term.
Server update. The general server optimisation (6 admits the closed-form solution (Appendix B):
N
m P
Mis
VU
N (1 (1+ + m)(m)))
n1
+Σ p(mo
mi me,p) (18)
where p(mo. m,, p) = pm - -pmom: - -pm;mp T. + +m Note that m, 's are fixed from clients' latest variational parameters. It is interesting to that mi in (18) generalises the
aggregation step of averaging local models in Fed-Avg (McMahan et al. (017) and related methods: when p 1 (no dropout), it almos equals client model averaging. Also, since p(m), m,, p = =) = (m, - (:-mmmmm -0m) when p I, Vi essentially estimates the sample scatter matrix with (N +1 samples, namely clients' m.'s s and server's prior 140 = O, measuring how much they deviate from the center mij The dropout is known to help regularise the model and lead to better generalisation (Gal & Ghahramani] 2016), and with p < 1 our forms a principled optimal solution.
Global prediction. The inner integral of (8) becomes the multivariate Student-t distribution. Then the predictive distribution for new test input x can be estimated as
(lo 1)Vo S Eppurle g(t)), where g(s) bng- -d11 0; mo, lo(ng - d 1) ) (19)
S
p(y |x", DI.N)= =
where (a, B) is the multivariate Student-f with location a, scale matrix B, and d.o.f.
Personalisation. With the given personalisation training data DP, we follow the general framework
in (11) to find v0) R p(0) DP, (²) ) in a variational way, where o² obtained from (34). We adopt the same spiky mixture form 16 for v()), which leads to the learning objective similar to (7).
3.2 MIXTURE MODEL
Our motivation for mixture is to make the prior p(0, o) more flexible by having multiple different
prototypes, diverse enough to cover the heterogeneity in data distributions across clients. We consider:
К
1 p()) II (, I), p(Ol))= - Σ K N(B:
(20)
where @ = {MI...... m contains K networks (prototypes) that can broadly cover the clients data distributions, and is the hyperparameter that captures perturbation scale, chosen by users or learned from data. Note that we put equal mixing proportions 1/K due to the symmetry, a priori. That is, each client can take any of (tj 's equally likely a priori. For the variational densities, define:
K
qs (0.) = (0,,m, II),
( = II (i
(21)
j=1
where (r (Lo) and my (Lx) are the variational parameters, and small constant (e.g., 10 10-4). Client update. For our model choice, the general client update (5) reduces to (details in Appendix C):
K
||m - -1y
min Ey [x,( log p( D.10.) - log
exp
(22)
20
It is interesting to see that 22 can be seen generalisation of proximal regularisation term in Fed-Pros is extended to multiple global models Tú 's, penalizing the local model (m,, straying away from these prototypes. And if we use single prototype (K = 1), the optimisation (22) exactly reduces to the local update objective of Fed-Prox. Since log-sum-est approximately equal to max, the regularisation term in (22) effectively focuses on the closest global prototype ry from the current local model m4, which is intuitively well aligned with our motivation.
(Li ct al. 2018), where the
is
'Only the constant added to the denominator, which comes from the prior and has the regularising effect. 'In practice use a single sample (S 1) for computational efficiency.
5

---

Under review as a conference paper at ICLR 2023
Server update. The general form
66 can be approximately turned into (Appendix C for derivations):
K
min 1
Σ
log
exp
|m=-ryl
(23)
tra
2
2a²
j=1
i=1
Interestingly, (3 generalises the well-known aggregation step of averaging local models in Fed-Avg and related methods: Especially when K = 1, (23) reduces to quadratic optimisation, admitting the optimal solution rf N+o² Σx= m,, The extra term o² can be explained by incorporating an extra zero local model originating from the prior (interpreted as neutral model) with the discounted weight o² rather than Although (23 for > can be solved by standard gradient descent, we apply the pnpectation-Maximisttoonn (EM) algorithmf (
(Dempster et al. 1977) instead:
(M-step) r = E c(jli) m,
(E-step) e(ili)
(24)
Σソ= e
N' + A I0=)
d(ib)
The M-step (server update) has intuitive meaning that the new prototype ry becomes the weighted average of the local models m,'s ' where the weights c(jlt) are determined by the proximity between my and ry (i.e., those my's that are closer to ry have more contribution, and vice versa). This can be seen as an extension of the aggregation step in Fed-Avg to the multiple prototype case.
Global prediction. We slightly modify our general approach to make individual client data domi- nantly explained by the most relevant model r,, by introducing a gating function from the mixture of experts (Jacobs et al. 1991 Jordan & Jacobs 1994). See Appendix C for details.
Personalisation. With v(O) of the same form as q;()),. the VI learning becomes similar to (2))
4 THEORETICAL ANALYSIS
We provide two theoretical results for our Bayesian FL algorithm: (Convergence analysis) As a special block-coordinate optimisation algorithm, we show that it converges to an (local) optimum of the training objective 33: (Generalisation error bound) We theoretically show how well this optimal model trained on empirical data performs on unseen test data points. Due to space limit, full details and proofs are described in Appendix DIE and we only state the thcorems and remarks here.
Theorem 4.1 (Convergence analysis). We denote the objective function in 3 by f(x) where x [wo, I1, INN] NN] corresponding to the variational parameters To := Lo, Iı Li. IN == Ln- Let nh = for some constant L. and チ" = = ET= x', where t is the batch iteration counter, x' is the iterate at by following FL algorithm, and N (S N) is the number of participating clients at each round. With Assumptions 1-3 in Appendix D the following holds for any T:
N N+N, .T+I D2 N₂ where x² is the (local) optimum, D, and R2 are some constants, and the expectation is taken over randomness in minibatches and selection of participating clients.
E[S(+²)) - [(x'
-00-10 VT
(25)
T
Remark. It says that 3² converges to the optimal point x" in expectation at the rate of O(1/Vi, This rate asymptotically equals that of the conventional non-block-coordinntt holistic) SGD algorithm.
Theorem 4.2 (Generalisation eтror bound). Assume that the variational density family for q:(0,) Po.,, Pi) be the expected squared Hellinger distance between the true class distribution Pi Pi(ylx) and model's 's Pa, (y(a:) for client i's data. The optimal solution (0.)) (0)
rich enough subsume Gaussian. Let d (Po,,
of the optimisation problem (3) satisfies:
1
F. co (Po,,
PP.,,' P') S o (H)
+ C C²
Fa Pa+a
Σ1
N
x
(26)
N
N
with high probability, where C,C² O are constant, X; mingge Il - 1'1 is the best error within our backbone network family θ and rw: Cn - 0 as the training data size n. - Co.
Remark. It implies that the optimal solution of 33 (attainable by our cloccccoordinate FL algorithm)
is asymptotically optimal, since the RHS of (26 converges to 0 as the training data size n - Do.
Instead of performing several EM steps until convergence, in practice find only one F.M step sufficient.

---

Under review as a conference paper at ICLR 2023
Global prediction (5- 1000
Pettora (s 100)
Global prodiction 1s- 1001
Personaiiaaioon (s
415
2450
623
40.6
20
Preset (Rooelev))
Fed BA2U
Feset (hoselie))
Ted-BAZU
N.W [Cura)
11 (0una)
Erceessee (haseeee
20
Entertiee (Facelin))
10 Eny
10
10
19
Cate namber nenwerks)
(a) Sensitivity Pdrop in NIW
(b) Our Mix. vs. ensemble baselines (K varied)
Figure 2: Hyperparameter sensitivity analysis and comparison with simple ensemble baselines.
RELATED WORK
General FL approaches. FedAvg (MMMahan et al. 2017 is the pioneering seminal work on FL, which proposed fairly intuitive local training and global aggregation strategies with minimal training and communication complexity. A potential issue of divergence of global and local models due to the separated steps of local training and aggregation was addressed by model regularisation in the follow- up works (Li et al. 2018 Acar et al. 202T), which is shown to help the global model converge more reliably. Recent approaches aimed to incorporate benefits from existing machine learning approaches including domain
clustering, multi-tat learning, transfer learning, and meta-learning. Due to the lack of space, we leave related references in Appendix !|
Comparison to existing Bayesian FL approaches. Some recent studies tried to address the FL problem using Bayesian methods. As we mentioned earlier, the key difference is that these methods do not introduce Bayesian hierarchy, and ultimately treat network weights 0 as random variable shared across all clients, while our approach assigns individual O, to each client i governed by a common prior p(). The nor-beeralllical approaches must all resort to ad hoc heuristics strong assumptions in parts of their algorithm. Due to the lack of space, we leave related references and discussions in Appendix I
EVALUATION
In this section we evaluate the proposed hierarchical Bayesian models on two benchmark datasets: the popular CIFAR-100 and the challenging corrupted version (CIFAR-C-100) that renders the client data more heterogeneous both in input images and class distributions.
FL settings for CIFAR-100. We follow the settings similar to those used in (Oh al. 2022); in particular the client data distributions are heterogeneous non-iid, formed by the sharding- based class sampling (McMahan et al., 2017). More specifically, we partition data instances in each class into non-overlapping equal-sized shards, and assign s randomly sampled shards (over all classes) to each of N clients. Thus the number of shards per user s can control the degree of data heterogeneity: small leads to more heterogeneity, and vice versa. The number of clients N = 100 (each having 500 training. 100 test samples), and we denote by f the fraction of participating clients. So, N,- = |N. clients are randomly sampled at each round to participate in training. Smaller / makes the FL more challenging, and we test two settings: ʃ f= 1.0 and 0. Lastly, the number of epochs for client local update at each round is denoted by 7 where we test 7 = 1 and 10, and the number of total rounds is determined by т as 3200T] for fairness. Note that smaller 7 incurs more communication cost but often leads to higher accuracy.
FL settings for TIFAR-100-Coree 2019) makes CIFAR-100' test split (10K images) corrupted by 19 different noise processes (e.g.. Gaussian, motion blur, JPEG). For each corruption type, there are 5 corruption levels, and we use the severest one. Randomly chosen 10 corruption types are used for training (fixed) and the rest 9 types for personalisation. We divide N = 100 clients into 10 groups, each group assigned one of the 10 training corruption types exclusively (denoted by D the corrupted data for the group e = 1,.... 10). Each D is partitioned into 90% 10% training/test splits, and clients in each group (LN/10) clients) gets non-iid train/test subsets from DS's train/test splits by following the sharding strategy with = 100 or 50. This way, the clients in different groups have considerable distribution shift in input images, while there also exists heterogeneity in class distributions even within the same groups. For the FL-trained models, we evaluate global prediction two datasets: clients' test splits from the 10 training corruption types and the original (uncorrupted) CIFAR's s training split (50K images). For personalisation, partition the clients into 9 groups, and assign one of the 9 corruption types to each group exclusively. Within each group we form non-iid marding-based subsets similarly, and again split the data into the 90% training/finetningg split and 10% test. Note that this personalisation
(CIFAR-C-100), The dataset (Hendrycks & Dietterich
7

---

Under review as a conference paper at ICLR 2023
Table 1: (CIFAR-100) Global prediction and personalisation accuracy.
(a) Global prediction performance (initial accuracy)
FL settings
Our Methods
Fed-BA Fed-BABU
Fed-Avg
Fed-Prox
pFedBayes
7
NIW
Mix. (K = 2)
1
49.76 19.76+0.12
49.37=0.30
42.35 12.35
40.87=0.62
41.49
37 37.23#00.8
0.1
10
29.02
29.02
27. t7.93
28. 28.26 20.19
27. 27.11 t0.11
28.21 11142
100
1
57. 17.80
52, 52.94 94
48, 18.17
47. 14 20 4..00..00
47 ...49
44.89+00.32
1.0
10
29. 29.53 #0.42
10.15
28, 28.67700.51
28. 79.6
27.43
28. 28.25
37.54 20.25
38.07=20.0
35.04 #0.56
27. 27.48 20.86
34.73t00.21
31.49440.18
18.99#0.03 10
18.95 18.95-40.13
18.54
14.69
16.84
17.93:10.68
10
50.40+0.11
49.52##0.8
45.41#0.11
37.10
44.33
$9.95++0.89
1.0
10
22.87#0.41
23.59 #0.47
21.92 +0.66
17.38 +0.32
19.54
21 21.85 21.85+t0.00 +0.00
(b) Personalisation performance
FL settings
Our Methods
Fed-BABU
Fed-Avg
Fed-Prox
pFedBayes
7
NIW
Mix. (K = 2)
1
54.16
56.17#0.16
50.43
16.43 46.43
49.9 49.91 +0.78
15.83=t1.12
0.1
10
36.38
36.32#0.27
35.15
33.57
33. 23.222
35. 74:11666
100
1
60.86
58.82
$5.87
55. 55.50
13.00
1.0
35.92* 10
36.22
35.58 40.24
33.82+11.00
33.70
35.57++1002
79.41 2:024
79.70#0.19
75.44
7..36
75. 75.06
73.93-10.14
0.1
67.35 10
67.35
67.57=0..2
06.24400.8
61.39
64.86
65.82:10.33
10
1
82.71
+0.37
81.03
78.92#t0.22
76.98+t0.66
76.98
6.666
+0.56
18.088+0.2
1.0
10
67.78
+1.02
66.74 20.27
66.25 0.46
63.81 4:040
63.81 t0.51
66.15 1..29
setting is more challenging compared to CIFAR-100 since the data for personalisation are utterly unseen during the FL training stage. We test T = 1 and 4 scenarios.
Experimental settings. Our implementation is based on Oh et al. 2022) where we use the Mo- bileNet (Howard et al., (017)) as backbone, and follow the body-update strategy: the classification head (the last layer) randomly initialised and fixed during training, with only the network body updated (and both the body and head updated at the personalisation stage). We report experimental results all based on this body-update strategy since we observe that it considerably outperforms the full update for our models and all other competing methods. The hyperparameters in our models are: (NIW) e = 10-4 and = 0.001 (See below for ablation study other values); (Mixture) = 0.1, e 10-4 mixture order K = 2 (See Appendix F.1 for results with other values), and the gating network has the same architecture as the main backbone, but with output cardinality changed to K. The other hyperparameters including batch size (50), learning rate (0.1 initially and decayed by 0.1) and the number of epochs in personalisation (5), are the same as those in (Oh et al. 2022).
Main results. In Table D (CIFAR-100) and Table (CIFAR-C-100), (NIW and Mixture with K
compare methods 2) against the popular ar FL methods, Fed-BABU (Oh et al.) (2022).
Fed-Avg (McMahan al. 2017). Fed-Prox (Li et al. 2018). and the recent PFedBayes (Zhang et 2022). The latter is especially interesting to contrast with as it is based on variational inference, most closely related to ours. We run the competing methods (implementation based their public codes) with default hyperparameters (e.g., = 0.01 for FedProx) and report the results. First of all, our two models (NIW and Mix.) consistently perform the best (by large margins most of the time) in terms of both global prediction and personalisation for nearly all FI. settings on the two datasets. This is attributed to the principled Bayesian modeling of the underlying FL. data generative process in our approaches that can be seen as rigorous generalisation and extension of the existing intuitive algorithms such as Fed-Avg and Fed-Prox. In particular, the superiority of our methods to the other Bayesian approach Baee verifies the effectiveness of modeling client-wise latent variables A, against the commonly used shared modeling, especially for the scenarios of significant client data heterogeneity (e.g., CIFAR-C-100 personalisation on data with unseen corruption types).
(Ablation) Hyperparameter sensitivity. We test sensitivity to some key hyperparameters in our models. For NIW, we have p= Pdrop- the MC-dropout probabbiity where we used parop 0.001
8

---

Under review as a conference paper at ICLR 2023
Table 2: (CIFAR-C- 100) Global prediction and personalisation accuracy.
(a) Global prediction (initial accuracy) on test splits for the training corruption types
FL settings
Our Methods
Fed-BABU
Fed-Avg
Fed-Prox
PFedBayes
7
NIW
Mix. (K 2)
1
81.22#0.14
80.34
79.45=0.71
70.01 70.01 #0.77
18.94#0.82
71.14
0.1
+7.74 4
63.81 20.84
63.58=1.25
18.47 1..771..2 1.26
60.96 1.11
57.88 11.61
100
1
91.20 91.26*00.8 +0.83
86.84+00.22
86, 86.84 84
85. 10
87. 07.033
82,80
1.0
73.58=1.0 1.02
+4.25
71.00 71.03#0.77
58.32#0.007
65.76 65.76+00.10
10010
64.07
78.63=0.5
79.36=t0.24
17.44
68.27#0.63
78.31 #0.93
67 67.95 07.25-20.22
0.1
65.08 4
63.52
62.65-#0.12
43.57-11.99
48.70#11.89
$4.300-:1.30
50
88.24+0.71
86.44=0.97
#0.67
86.00
82.32+0.37
1.0
4
70.33 00118
70.19
67. 67.66*0.75
52.48 1.06
60.72 11
60.17
(b) Global prediction (initial accuracy) on the original (uncorrupted) CIFAR-I00 training sets
FL settings
Our Methods
Fed-BABU
Fed-Avg
Fed-Prox
pFedBayes
f
NIW
Mix. (K 2)
1
+0.11
36.99
31.76=0.60
34.40=00331
35.44
35.78 t0.73
0.1
30.84+0.07
30.60
28.31 #0.28
20.70
28.09#0077
29.12
100
1
3..32200.2
38.35 40.73
35.58
36.34 +0.14
36.32 32.00
37.37
1.0
4
30.67#0.12
30.40
28.60
28.60
30.31 #0.87
27.95
29.14 +0.57
41.04
36.41 10044
25.4120.55
34.13 20.50 5.50
2..37200.50 50
35.6820.27
0.1
4
32. 32.29
31.50 31.50=20.34
29.68
29.19 14
29.20 29.20
30.10
1
11..44
38.54 38.54+0042 +0.42
36.09
36.09-0.22
35.78 70.788...8
37.13 60.65
38.39 +0.21
1.0
4
32.17
40.48
30.68 30.68 40.46
29.28 #017
30.45 40.44
28.73 18.733t0.20 +0.26
29.74 20.25
(c) Personalisation performance on the 9 held-out corruption types
FL. settings
Our Methods
Fed-BABU
Fed-Avg
Fed-Prox
PFedBayes
7
NIW
Mix.
1
72.63=2.15 2.13
74.16#3.00
69.93=1.24
62.32 1.46
72.94
62.5224.01
0.1
62.74 62.74 4
61.5611.60
60.33=2.12
53.35
56.91#1.40
53.71
100
83.62 1.84
80.88
77.55 #2.05
83.44 83.44+1.6s
+1.68
80.96 +3.18
72.44 72.44+0.13 +0.13
1.0
4
64.84* 1.05
67.35 #1.46
53.25 #2.19
53.34 1.22
41.39 11118
43.41 +1.92
1
75.50
67.3342.88
59.17=2.46
57. 57.7720.80
56.34 80
53.1740.51
0.1
1.29 4
46.39 #0.83
44.74=1.30 1.30
44.60 #2.36
11.94
37.24
50
1
81.46=u.0
b3.11
67.43* #2.68
72.52 1.02
71.52 +2.02
68.47+1143
1.0
50.84+0.74
40.90
40.87=3.01
10.10
41.95#1.28
41.27#0.14
in the main experiments. In Fig. 2(a) we report the performance of NIW for different values (Pdrop = 0, 10 1011))) on C0F with (s = 100, f = 0.1, = 1) setting. We see that the performance is not very sensitive to Parop unless is too large (e.g., 0.01). For the Mixture model, different mixture orders
2, 5, 10 are contrasted in Fig. (b). As seen, having more mixture components does harm (no overfitting). but do not further improvement over K 2 in our experiments (See also results on CIFAR-C-100 in Table Appendix).
Further analysis. In Appendix, provide further empirical results: (i) comparison between mixture model and simple ensemble baselines (Fig. bb) and F2 and (ii) actual running times (F3)
7 CONCLUSION
have proposed a novel hierarchical Bayesian approach to FL where the block-coordinate descent solution to the variational inference leads to viable algorithm for FL. Our method not only justifies the previous FL algorithms that look intuitive but theoretically less underpinned, but also generalises them even further via principled Bayesian approaches. With strong theoretical support in convergence rate and generalisation error, our approach is also empirically shown to be superior to recent FL approaches by large margin several benchmarks with various FL. settings.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Durmus Alp Emre Acar, Yue Zhao, Ramon Matas Navarro, Matthew Mattina, Paul N Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. International Conference on Learning Representations, 2021.
Idan Achituve, Aviv Navon, Yochai Yemini, Gal Chechik, and Ethan Fetaya. GP-Tree: A Gaussian Process Classifier for Few-Shot Incremental Learning. In International Conference on Machine Learning, 2021a.
Idan Achituve, Aviv Shamsian, Aviv Navon, Gal Chechik, and Ethan Fetaya. Personalized federated learning with Gaussian processes. Advances in Neural Information Processing Systems, 2021b.
Maruan Al-Shedivat, Jennifer Gillenwater, Eric Xing, and Afshin Rostamizadeh. Federated Learning via Posterior Averaging: A New Perspective and Practical Algorithms. International Conference on Learning Reprexentations, 2021.
Jincheng Bai, Qifan Song, and Guang Cheng. Efficient Variational Inference for Sparse Deep Learning with Theoretical Guarantee. In Advances Neural Information Processing Systems, 2020.
Amir Beck and Luba Tetruashvili. On the convergence of block coordinate descent type methods. SIAM Journal on Optimization, 23(4):2037-2060, 2013.
Dimitri P. Bertsekas. Nonlinear Programming. Athena Scientific, 2016.
Christopher M. Bishop. Pattern Recognition and Machine Learning. Springer, 2006.
D. M. Blei, A. Kucukelbir, and J. D. McAuliffe. Variational inference: A review for statisticians. Journal of the American Statistical Association, 112(518):559-877, 2017.
S. Boucheron, G. Lugosi, and P. Massart. Concentration Inequalities: Nonasymptotic Theory of Independence. Oxford University Press, 2013.
Michael Braun and Jon McAuliffe. Variational inference for large-scale models of discrete choice. arXiv preprint arXiv:0712.2526. 2008.
Christopher Briggs, Zhong Fan, and Peter Andras. Federated learning with hierarchical clustering of local updates to improve training non-iid data. International Joint Conference Neural Networks, 2020.
Fei Chen, Mi Luo, Zhenhua Dong, Zhenguo Li, and Xiuqiang He. Federated meta-learning with fast convergence and efficient communication. arXiv preprint arXiv 077766 2018.
Hong-You Chen and Wei-Lun Chao. FedBE: Making Bayesian Model Ensemble Applicable to
Federated Learning. In International Conference on Learning Representations, 2021.
Yiqiang Chen, Xin Qin, Jindong Wang, Chaohui Yu, and Wen Gao. Fedhealth: A federated transfer learning framework for wearable healthcare. IEEE Intelligent Systems, 35(4):83-93, 2020.
Badr-Eddine Abdeief-Abdellatif.. Convergence Rates of Variational Inference in Sparse Deep Learning, 2020. International Conference Machine Learning.
A. P. Dempster, N. M. Laird, and D. B. Rubin. Maximum likelihood from incomplete data via the algorithm. Journal of the Royal Statistical Society. Series B (Methodological), 39(1):1-38, 1977.
Canh T Dinh, Nguyen Tran, and Tuan Dung Nguyen. Personalized federated learning with Moreau envelopes. Advances in Neural Information Processing Systems, 33, 2020.
Canh Dinh, Tung T Vu, Nguyen H Tran, Minh N Dao, and Hongyu Zhang. Fedu: A uni- fied framework for federated multi-task learning with Laplacian regularization. arXiv preprint arXiv:2102.07148, 2021.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning: A meta- learning approach. arXiv preprint arXiv:2002.07948, 2020.
10

---

Under review as a conference paper at ICLR 2023
Yarin Gal and Zoubin Ghahramani. Dropout as Bayesian Approximation: Representing Model Uncertainty in Deep Learning. International Conference on Machine Learning, 2016.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. In International Conference on Leurning Representations, 2019.
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. arXiv preprint arXiv: 1704.04861, 2017.
R. A. Jacobs, M. I. Jordan, S. Nowlan, and G. E. Hinton. Adaptive mixtures of local experts. Neural Computationn 37, 1991.
M. I. Jordan and R. A. Jacobs. Hierarchical mixtures of experts and the EM algorithm. Neural Computation, ((2):881-214 1994.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J. Reddi, Sebastian U. Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic Controlled Averaging for Federated Learning. arXiv preprini arXiv: 1XXix:991000638: 2019.
D. P. Kingma and M. Welling. Auto-encoding Variational Bayes. In International Conference on Learning Representations, 2014.
Nikita Kotelevskii, Maxime Vono, Eric Moulines, and Alain Durmus. FedPop: A Bayesian Approach for Personalised Federated Learning. In Advances in Neural Information Processing Systems, 2022.
John Langford and Rich Caruana. (Not) Bounding the True Error. In Advances in Neural Information Processing Systems, 2001.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated Optimization in Heterogeneous Networks. arXiv preprint arXiv:1812.06127, 2018.
Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated
learning through personalization. International Conference on Machine Learning, 2021.
Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang. On the Convergence of
FedAvg on Non-IID Data. arXiv preprint arXiv:1907. 02189, 2019.
Xiang Li, Kaixuan Huang, Wenhao Yang. Shusen Wang, and Zhihua Zhang. On the convergence of fedavg on non-iid data. In International Conference on Learning Representations, 2020.
Wesley Maddox, Timur Garipov, Pavel Izmailov, Dmitry Vetrov, and Andrew Gordon Wilson. A Simple Baseline for Bayesian Uncertainty in Deep Learning. In Advances in Neural Information Processing Systems, 2019.
Yishay Mansour, Mchryar Mohri, Jac Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning. arXiv preprint arXiv:2002. 10619, 2020.
Othmane Marfoq, Giovanni Neglia, Aurélien Bellet, Laetitia Kameni, and Richard Vidal. Federated Multi-Task Learning under a Mixture of Distributions. In Advances in Neural Information Processing Systems, 2021.
Andreas Maurer. Note on the PAC Bayesian Theorem. arXiv preprint arXiv:0411099, 2004.
David McAllester. Some pac-bayesian theorems. Machine Learning, 37:355-363, 1999.
H. Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera y Arcas. Communication-Efficci Learning of Deep Networks from Decentralized Data, 2017. Al and
Statistics (AISTATS).
Jachoon Oh, Sangmook Kim, and Se- Young Yun. FedBABU: Towards Enhanced Representation for Federated Image Classification. In International Conference Learning Representations, 2022.

---

Under review as a conference paper at ICLR 2023
D. Pati, A. Bhattacharya, and Y. Yang. On the Statistical Optimality of Variational Bayes, 2018. AI and Statistics (AISTATS).
Daniel Peterson, Pallika Kanani, and Virendra J. Marathe. Private Federated Learning with Domain Adaptation. arXiv preprint arXiv:19/2.06733, 2019.
Omar Rivasplata, Vikram M Tankasali, and Csaba Szepesvari. PAC-Bayes with Backprop. arXiv preprint arXiv:/908. 100x::990..07380, 2019.
Matthias Sceger. Pac-bayesian generalization error bounds for gaussian process classification. Journal of Machine Learning Research, 3:233-269, 2002.
Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet Talwalkar. Federated multi-task learning. arXiv preprint arXiv: /705. 10467, 2017.
Benyuan Sun, Hongxing Huo, Yi Yang, and Bai. Partialfed: Cross-domain personalized federated learning via partial initialization. Advances in Neural Information Processing Systems, 34, 2021.
Niklas Thiemann, Christian Igel, Olivier Wintenberger, and Yevgeny Seldin. A strongly quasiconvex PAC-Bayesian bound. International Conference on Algorithmic Learning Theory, 2017.
Huahua Wang and Arindam Banerjee. Randomized Block Coordinate Descent for Online and Stochastic Optimization. arXiv preprint arXiv: 1407.0107, 2014.
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H. Vincent Poor. Tackling the Objective In- consistency Problem in Heterogeneous Federated Optimization. arXiv preprint arXiv:2007.07481, 2020.
Stephen Wright. Coordinate descent algorithms. Mathematical Programming, 151(():3-34,2)) 2015.
Hongwei Yang, Hui He, Weizhe Zhang, and Xiaochun Cao. Fedsteg: A federated transfer learning framework for secure image steganalysis. IEEE Transactions on Network Science and Engineering, 8(2):1084-1044, 2020.
Liling Zhang, Xinyu Lei, Yichun Shi, Hongyu Huang, and Chao Chen. Federated Learning with Domain Generalization. arXiv preprint arXiv:2111.10487, 2021.
Xu Zhang, Yinchuan Li, Wenpeng Li, Kaiyang Guo, and Yunfeng Shao. Personalized Federated Learning via Variational Bayesian Inference, 2022. International Conference on Machine Learning.
12

---

Under review as a conference paper at ICLR 2023
A ELBO DERIVATION FOR GENERAL FRAMEWORK
We derive the ELBO objective for the general Bayesian FL framework.
g())= II. (0.) p(DI.N)
(l. 01.) p(o, 01.N PI:N)) Fyg log p()) II, (0, (D:
(27)
KL(g())lp(O)) + Σ (Bn(a)l log p( (D. B()) (B)))()))lpp(..))))))
+ log pl DI:N).
(28)
Since KL divergence is non-negative, -L(L) must be lower bound of the data log- log-likelihodd log p( DI.N), rendering (LL) as our objective function (to be minimised).
B NORMAL-INVERSE. WWW-IIEEESSE-SSHARRY (NIW) MODEL (DETAILED VERSION)
We define the prior as conjugate form of Gaussian and Normal-Inverse-V Wisharr.. More specifically, each local client has Gaussian prior p({,(o) = (0;; μr, Σ) where ø = (j,, Σ), and the global latent Sartaal-nnneee--iiiiit (NIW).
variable Φ is distributed as a conjugate prior which
p = NIW(H, A) (N; (o, ²) IW(E;; Eo, v0), p(0, x) (N,0; 14, Σ). i= I,..,,N. where A {Mo, Eu, Au } is the parameters of the NIW. Although A can be learned via data marginal likelihood maximisation (e.g., empirical Bayes), but for simplicity we leave it fixed as: 140 = 0, Σn
(29)
(30)
Ao and un = d+2 2 where d is the number of parameters in 0; or μ.. This choice ensures that the mean of E equals I, and is distributed as zero-mean Gaussian with covariance Σ.
Next, our choice of the variational density family for q()) the NIW, not just because it the most popular parametric family for a pair of mean vector and covariance matrix Φ = (μ, )), but it can also admit closed-form expressions in the ELBO function due to the conjugacy as derive in Sec. B. ((6)= ;= N (ZW( {mo, Vo. lo. no N N(u; mo, 1o -)) IW(E; Vo, no).
(31)
Although the scalar parameters on can be optimised together with mo. Vo, their impact is less influential and we find that they make the ELBO optimisation a little bit cumbersome. So we aim to estimate their optimal values in advance with reasonably good quality. To this end, we exploit the conjugacy of the NIW prior-posterior under the Gaussian likelihood. For each 0;, we pretend that we have instance-wis representative estimates 0: (z. y), one for each (x., v) E D,. For instance, one can view 0,(x; y) as the network parameters optimised with the single training instance (xx y). Then this amounts to observing [D] (= EN [D,)) Gaussian samples (z y) ~ (0,;;; J.. Σ) for (z, ~ D, and i = = N. Then applying the NIW conjugacy, the posterior is the NIW with =
to = Ao = estimates for the optimal lo. ng. and we fix them throughout the variational optimisation. Note that this is only heuristics for estimating the scalar parameters lg. no quickly, and the parameters mo. Vo are determined by the principled ELBO optimisation (Sec. B. That is, Lo {mo. Vo). Since the dimension d is large (the number of neural network parameters), we restrict Vo to be diagonal for computational tractability.
and no = Po + | = [D] + d + 2. This gives us good approximate
The density family for q: (0.))s can be a Gaussian, but we find that it is computationally more attractive and numerically more stable to adopt the mixture of two spiky Gaussians that leads to the MC-Dropout (Gal & Ghahramani] 2016). That is,
q: (0,) = II (p (0, m1 21) + -11-p)-N((00:000 -V(8.000,221)
(32)
where (i) m. is the only variational parameters (L, {m+)), (ii -[!] indicates the specific colum- nllayer in neural network parameters where goes over layers and columns weight matrices, (iii) p is the (user-sper specifed)) hyperparameter where 1 corresponds to the dropout probability, and (iv) is a tiny constant (e.g.. 10 6) that makes two Gaussians spiky, close to the delta function. Now we provide more detailed derivations for the client optimisation and server optimisation.
13

---

Under review as a conference paper at ICLR 2023
B.1 DETAILED DERIVATIONS FOR NIW MODEL
Client update. We work on the objective function in the general client update optimisation (5)- We note that g(0) is spiky since our pre-estimated NIW parameters lo and no are large (as the entire training data size [D] is added to the initial prior parameters). Due to the spiky (()), we can
accurately approximate the second term in (5) as:
By(d) [L(( |pp00 [ s (L(g: (L))((:))ll())))))))
(33)
where o² (1x", Σ²) is the mode of (()), which has closed forms for the NIW distribution:
Vo
mo: E
(34)
To + +
In 33 we have the KL divergence between mixture of Gaussians 32 and Gaussian (30 Similar
to (Gal& & Ghahramani 2016). we apply the approximation KL a. Rs Σ. ag (L
well as the reparametrised sampling for (32h. which allows us to rewrite ( as:
min (m,) == log (,)) 2
(no + d 1)(m; - (m)) Vo (m; - mo), m,
(35)
where my is the dropout version of m,, i..., reparametrised sample from 32 Also, use minibatch version of the first term for a tractable SGD update, which amounts to replacing the first term by the batch average (((e.)))Baatl factor of 1/ D.I. Note that mo and Vo are fixed during the optimisation. Interestingly (35) generalises the famous Fed-Avg (McMahan et al. 2017) and Fed-Prox (Li et al.
log p(y) [z. m.)] while downweighing the second term by the
2018): With p = (i.e., no
dropout) and setting Vo= al for some constant a, formula for Fed-Prox where a controls the impact of the proximal term.
that 35 reduces to the client update
Server update. The server optimisation (6) involves two terms, both of which we will show admit closed-form expressions thanks to the conjugacy. Furthermore, we show that the optimal solution (ma; Vo) of (6) has an analytic form. First, the KL term in (6) is decomposed as:
KL (IW((; Vo, No). Eg. v)) + EIW(I:Vo.a))
KL (N(() mg. 1o" E)I Po, Ap [)]]
(36)
By some algebra, (36 becomes identical to the following, up to constant, removing those terms that are not dependent on mo. (See Appendix 2 for derivations):
No Tr(oo
so log Vol + Aono(uo -ma) VVo 'wo mo)
(37)
Next, the second term of 66 also admits closed form as follows (Appendix B.2 for details):
E Eg(0)) (0.) [o,g((; = 2
no
V pmm 3 m, pmg Vo mm;- pm, Vo mo + mg Vo
mg
1 ni
+ log Vol + TT(((V + const.
(38)
That is, server's loss function Co is the sum of (37) and 38) We can take the gradients of the loss with respect to ma, Vo as follows (also plugging Ho O. Eo
I, Au 1, v/ = d 2):
OCo Omn
noVo
(N + 1)mo - Σ m
(39)
aCo
1
N no N (N + d 2) + nomomi + +no Σ
p(mo, mi m
(40)
Ov
2
where p( mi; p) = pmmm, pmmm, pmn.m + mom T
We set the gradients to zero and solve for them, which yields the optimal solution:
N
mi = N+1 11
m.
Ve
no
(1+Nee) mm +
Elmm m,, p)
N
(41)
Note that my's 's are fixed from clients' latest variational parameters.
It is interesting to see that mā in (4] generalises the well-known aggregation step of averaging local models in Fed-Avg (McMahan et al., 2017) and related methods: when p = (i.e., no dropout), (irn-mm)) mm²) when almos equals client model averaging. Also, since p(mng mi, 1) (m, (mn;-m5)(mm -m))(m
p = 1, we can see that VO in 41 essentially estimates the sample scatter matrix with (N+1) +1) samples,
'Only the constant added the denominator, which comes from the prior and has the regularising effect.
14

---

Under review as a conference paper at ICLR 2023
namely clients' m,'s and server's prior fi 0, measuring how much they deviate from the center mp. is known that the dropout can help regularise the model and lead to better generalisation (Gal & Ghahramani
2016), and with p < 1 our (41) forms principled optimal solution.
Global prediction. In the inner integral of ( of the general predictive distribution, we plug
p(0 Φ) = (0; Σ) and NIW (()) of (31) This leads to the multivariate Student-f distribution:
(lo 1)Vo
p(0 ( dø
N/ E) NIW(()) do tmo d+1 0: mo: lo d 1)
(42)
where (a, B) is the multivariate Student-t with location a, scale matrix b, and d.o.f. Then the predictive distribution for a new test input x* can be estimated a
/ p(u"|o" 10",)) Imo -do1 0; mo: Io(no - d
(lo 1)Vo
13"10" DI....,DN) =
d0
(43)
s S Σ p(u" 1,",)())), where g(v) ~ Imo d+1 0: mo, lo(no - d + 1)
(lo 1)
(44)
Personalisation. With the given personalisation training data D we follow the general framework in (11) to find v(0) = p(8) Dr, ,d)) in a variational way, where o² obtained from (34). For the density family for v(0) we adopt the same spiky mixture form as (32).
v (0) - II (p m + (1 p) ((011),,
(45)
where m is the variational parameters. This leads to the MC-dropout-lic learning objective, min - p(" + - (no + d 1)(m mu) VV (m mu),
(46)
2
Once e is trained, our predictive distribution follows the MC sampling 12)
B.2 MATHEMATICAL DETAILS
The server optimisation (6) in our NIW model involves terms, both of which will show admit closed-form expressions thanks to the conjugacy. Furthermore, we show that the optimal solution (mo, Vo) of (6) has an analytic form. First, the KL term in (6) decomposed as: (I(((())lp(())
(+) = (((((()))())) ((p())((E))
(47)
=Ba()) log ({E)
(M
+ (E) (
log p(u [E]
(48)
p(E)
(EW Vo, no) IW Ea, Mb)) + E EIW(EEVVo.O) KLLN(p; mo 1' Σ)ソ
=ible
=:kl,
(49)
First we work on kla = EZWEE.V. [log IW(E; Vo, mo)] - EIWCE:Vo. flo IW( IoVE;; Eo, vo)]. From the definition Inverse- Wishart (assuming Σ = (d x d)),
d +
vd
log IW VV log |
- log | Tr( log log
log 2,
2
2
(50)
where Ta(-) is the multivariate Gamma function. We use the following facts from (Bishop
2006
Braun & McAuliffe 2008):
EIW(E log Σ = - -d 2 + log | ((v
(51)
(52)
where (-) is the digamma function. Applying these to the terms in kla yields:
kl. = 2 (Too + un log Vol) + const (w.rt, ma, Vo).
(53)
"I practice use single sample (S 1) for computational efficiency.
15

---

Under review as a conference paper at ICLR 2023
Next, using the closed- form expression for the KL between Gaussians, klg becomes:
klp EIW(E:VO.m)) [No - mm)) - - (mo - -mo]] + const (w.rt. mp, Vo)
Σ-
(54)
Aona
2 (mo mu) VV (NO ma) + const (w.r.t mu, Vo).
(55)
where in (55) use the fact (52) Combining 56) and (55). have:
(L((())lp()) = 2 no Tr VoVo + un log Vol + Aoooouo - mo) Vo - mo)) + const.
(56)
Next, we derive the second term of 6 ute stands for equality up to constant (w.ntt mo: Vo)).
1
Bg (d( (0,) llog [logo(0000)))) p(0))))] = E log
D + (0, 11) 1)/ 1-; -+²] + const (w.r.t. Ma, Vo)
(57)
Fute
EIW(I:Vo.
log |ل|
1 [0, - -4) - - ni]
(58)
2
2
1
Fute
log Vol
Tr BIW(ENVo,o))
Σ E- (N (I)),, [(0, - []
(59)
2
Fute
1
log Vol
Tr m))(E:
2 (p(mo, m.,)) 2 1/ 2;)
(60)
2
Futc
log Vol
Tr (p(mo: m. p). =e)onoooo
(61)
2
utc
110
2
pm Vo
T
Vo- 1m-m= pmg Vo- m, 'm "mm Vo-
'mo + mo 1o-1 mmg +
log Vol
70
+
TTVVV
(62)
where p(mo. m. pm.m pmom. pmmmm + mo and we use the definition of (6.)
in (22 and the fact 52
C MIXTURE MODEL (DETAILED VERSION)
Previously, the NIW model expresses our prior belief where each client acquires its own network parameters 0, a priori as a Gaussian- perturbed version of the shared parameters 14, namely 0, ~ N(,,)), 2). as in (14). This intuitively appealing, but may not be adequate for capturing more drastic diversity in local data across clients. In the situations where clients' local data distributions. as well as their domains and class label semanties, are highly heterogeneous (possibly even set up for adversarial purpose), it would be ideal to consider multiple different protolypes for the network parameters, diverse enough to cover the heterogeneity data distributions across clients. Motivated from this idea, introduce mixture prior model as follows.
First we consider that there are K network parameters (prototypes) that can broadly cover the clients data distributions. They are denoted high-level latent variables, , H }, and we let them distributed independently as standard normal a priori,
K
p
= II (, j=1
(63)
We here note some clear distinction from the NIW prior. Whereas the NIW prior [13) only controls the mean J and covariance Σ in the Gaussian, from which local models 0, are sampled, the mixture prior (63) is far more flexible in covering highly heterogeneous distributions. Fach local model is then assumed to be chosen from one of these K prototypes. Thus the prior distribution for 0 can be modeled as mixture,
K Σ
p(0
K N(O:H;011)
(64)
where o is the hyperparameter that captures perturbation scale, and can be chosen by users or learned from data. Note that we put equal mixing proportions 1/K due to the symmetry, a priori. That is, each client can take any of (1y 's equally likely priori.
16

---

Under review as a conference paper at ICLR 2023
We then describe our choice of the variational density (()) II, 9:(8,) approximate the posterior p(o, 01:N (I:N). First, (:(),) is chosen as a Gaussian,
g: (#,))) = (0;;; [(,,;,;e" 2
(65)
with small For (()) we consider Gaussian factorised over 24y but with small variances, that is,
K
( = II (u,;,,,))) _=1
(66)
where fry are variational parameters (Lg) and c small (e.g., 10-4). The main reason why we make (()) spiky that the resulting sear-deterministic q(()) allows for computationally efficient and accurate MC sampling during ELBO optimisation well as test time (global) prediction, avoiding difficult marginalisation (Sec. C.T for details). Although Bayesian inference in general encourages to retain many plausible latent states as possible under the given evidence (observed data), aim to model this uncertainty by having many (possibly redundant) prototypes jiy rather than imposing larger variance for single one (e.g., finite-sample approximation of smooth distribution).
C.I DETAILED DERIVATIONS FOR MIXTURE MODEL
With the full specification of the prior distribution and the variational density family, we are ready to dig into the client objective function (5) and the server
Client update. Since (()) is spiky, we can accurately approximate the second term of (5) as (L(q((0)) (p())))))) where 6² {ug = ry 1- i is the mode of g(Ø) since
(o ((gp ss Since (0.) is also spiky, (0,) Ip((, ())), the KL divergence between a Gaussian and a Gaussian mixture, can be approximated accurately using the single mode sample m, (6,), that is, (L, (p(((()))) s log (m,) - log p(m,lo")
p(0,
(67)
(68)
||m; - ry113
log Σ N(mmr o21) + const. log Σ exp
+ const. (69)
20
j=1
Note here that we use the fact that m, disappears log (m:)) following optimisation for client i:
Plugging it into 5 yields the
K
||m ryy
min (., log (D)) log Σ exp 1744
(70)
20
It is interesting see that 770 can be seen as generalisation Fed-Prox (Li 2018), where the proximal regularisation term in Fed-Prox is extended to multiple global models rý 's penalizing the local model (m,) straying away from these prototypes. And we use single prototype (K 1), the optimisation (70) exactly reduces to the local update objective of Fed-Prox. Since log sum exp approximately equal to max, the regularisation term (70) effectively focuses on the closest global prototype ry from the current local model m,, which is intuitively well aligned with initial modeling motivation, namely each local data distribution is explained by one the global prototypes. Lastly, also note that in the SGD optimisation setting where
only access minibatch B D; during the optimisation of (70). we follow the conventional practice: replacing the first term of the negative logelikeli by stochastic estimate ( B(a log p(yl x. 0,) and multiplying the second term of regularisation by -..
Server update. First, the KL term in 6 can be easily derived as:
K
(L(g())IP(()) = Σ
|r, const.
(71)
j=l
and the second term of (6) approximated as follows:
K
1
E Ey (ocop x (log ~ log
N(mar3y,07
(72)
K
j=1
K
[m - r3(12)
= log Σ
exp
+ const.
(73)
20
j=1
17

---

Under review as a conference paper at ICLR 2023
where the approximations in (72) bining the two terms leads to the optimisation problem for the server:
become accurate due to spiky q:(),) and q()), respectively. Com-
K
N
K
I|m, - r,ll
min (ry
Σ
Σ log
exp
(74)
2
20²
j=1
i=1 j=1
Interestingly, (74) generalises the well-known aggregation step of averaging local models in Fed- Avg (McMahan et al. (017) and related methods: Especially when K 1, (74) reduces to quadratic optimisation, admitting the optimal solution rf = N+a= Σx= my. The extra term o² in the denominator can be explained by incorporating an extra zero local model originating from the prior (interpreted neutral model) with the discounted weight o² rather than 1.
Although (74) for K > can be solved by standard gradient descent, the objective function resembles the (regularised) Gaussian mixture log- ligeleiehoood, and we can apply the Expectation- Motitionn-Miiiittttionnn (EM) algorithm (Dempster et al.) 1977) instead. Using Jensen's 's bound with convexity of the negative log function, we have the tollowing alternating steps
E-step: With the current {ry} each local model m,:
K fixed, compute the prototype assignment probabilities for
kig where kig = exp
||m, - -ryll3
c(il))
(75)
Σラ kis
20²
M-step: With the current assignments e(jli) fixed, we solve:
min 2/ Σ ||,111 + 20 {ry)
1
Σ c(jli) [|,, - ry11?
(76)
2
which admits the closed form solution:
DD= c() m,
j = 1,
(77)
+ E
The server update equation (77) has intuitive meaning that the new prototype ry becomes the weighted average of the local models m, 's where the weights c(jli) are determined by the proximity to ry (i.e., those m,'s 's that are closer to r, have more contribution, and vice versa). This can be seen as an extension of the aggregation step in Fed-Avg to the multiple prototype case.
Global prediction. By plugging the mixture prior p()) of (64) and the factorised spiky Gaussian (()) of (66) into the inner integral of (8) we have predictive distribution averaged equally over (ry}5= approximately, that is, J p(0 ()) g( de = K E5=1((r (=",rr)) Unfortunately this is not ideal for our original intention where only one specific model ry out of K candidates is dominantly responsible for the local data. To meet this intention, extend our model so that the input point x² can affect 0 together with Φ and with this modification our predictive probability can be derived as:
/ʃ ply" 0) p(0)=", Φ) D(.) ddds
plv" D
(78)
Rd // plu"|=", ()) p(0==", .,="))) .)) (()) deds
(79)
P( ,00 p(0 (ry) } de.
(80)
To deal with the tricky part of inferring p(0|z", fr, ) we introduce fairly practical strategy of fitting gating function. The idea is to regard p(0|=", {r)) as a mixture of experts (Jacobs et 1991: Jordan & Jacobs 1994) where the prototypes ry serving as experts,
p(0 x0 (1y
()((") 5(0 ry))
(81)
1=1
where 6(-) is the Dirac's s delta function, and g(x) is a gating function that outputs a K-dimensional softmax vector. Intuitively, the gating function determines which of the K prototypes {ry}5= the
Athough one can perform several EM steps until convergence, in practice, we find that only one EM step per round is sufficient.
18

---

Under review as a conference paper at ICLR 2023
model 0 for the test point x² belongs to. With (8 the predictive probability in 80 is written
ply" x=" DI:N)= = g p(y"|=""
(82)
However, since we do not have this oracle g(z), we introduce and fit a neural network to the local training data during the training stage. Let g(æ; (3) be the gating network with the parameters B. To train it, we follow the Fed-Avg strategy. In the client update stage at each round, while we update the local model my with minibatch B D; we also find the prototype closest to mi, namely j" := arg min; ||m ryll- Then we form another minibatch of samples {(x;) J=) (input x and class label. j"), and update g(æ; ()) by SGD. The updated (local) ( from the clients are then aggregated (by simple averaging) by the server, and distributed back to the clients as an initial iterate for the next round.
Personalisation. For porp(() Dr, Φ²) in the general framework 10 we define the variational distribution v(0) = D", (6) as:
v(0 = (0; m, 2T),
(83)
where is small positive constant, and is the only parameters that we learn. Our personalisation training amounts to ELBO optimisation for v(f) as in 11 which reduces to:
m - ry11?
min En(o)) log p( [P) log
exp
(84)
20
Once we have optimal m (i.c., v()), our predictive model becomes:
p(y xT, D", DD.N) s p(y"|=", m,),
(85)
which is done by feed-forwarding test input x² through the network deployed with the parameters m.
D CONVERGENCE ANALYSIS
Our (general) FL algorithm is a special cookdccord SGD optimisation of the ELBO function (3) with respect to the (N 1) parameter groups: Lo (of q(o: Lg)), L, (of qı (0): )),,, ..... ..... and L (of IN (ON; LN)), )). In this section we will provide a theorem that guarantees convergence of the algorithm to local minimum of the ELBO objective function under some mild assumptions. We will also analyse the convergence rate. Note that although our FL algorithm is special case of the general block-coordinate SGD optimisation, we may not directly apply the existing convergence results for the regular flock-coordinate SGD methods since they mostly rely on -vneraaaag blocks with cyclic or uniform random block selection strategies (Beck Tetruashvili] 2013 Wang Banerjee 2014). As the block selection strategy in our FL algorithm is unique with overlapping blocks and non-uniform random block selection, we provide our own analysis here. Promisingly, we show that in accordance with general regular block-coordinate SGD (cyclic/uniforn non-overlapping block selection), our FL. algorithm has O(1/V2 convergence rate, which is also asymptotically the same as that of the conventional (holistic, non-block-coordinaat SGD optimisation. Note that this section is about the convergence of our algorithm to an (local) optimum of the training objective (ELBO). The question of how well this optimal model trained on empirical data performs on the unseen data points will be discussed in Sec.
First we formally describe our FL algorithm as a block-coordinate SGD optimisation. For ease of exhibition, will simplify the notation: The objective function in (3) is denoted as f(x) where & = [Tn. T1, N] is the optimisation variables corresponding to Fo := Lo. X1 := TT. IN InN That is, Z0 is server's parameters while x; (i 1, 1,..,, N) is worker i's parameters. We let r. be the partial vector of x selected by the index set u S 1,. N). and x be the vector of the rest elements. Similarly (x. indicates the gradient vector with only elements at the indices in u. Let x' be the iterate at iteration f. Our FL algorithm is formally defined in Alg.
Now we state our convergence theorem. We first need the following mild assumptions:
Assumption 1. Our objective function f(xx) is locally convex, and f((z) is also locally convex for all iterations t, where fo is the minibatch version of defined on the minibatch data batch, (so that Epatch, [f (x)] = f())). Actually, the latter implies the former. Although the negative ELBO is in
'We also follow the Fed-BABU (Oh et al. 2022 strategy by updating only the body of (3 and fixing/sharing the random classification head across the server and clients.
19

---

Under review as a conference paper at ICLR 2023
Algorithm Bayesian FL Algorithm Block-Coordinate Descent.
We define the following hyperparameters:
N (s N) = the number of participating clients at each round.
M = the number of SGD iterations per round for updating the (participating) clients.
Let M, the number of SGD iterations per round updating server.
Let 7) = the reciprocal
SGD learning rate at iteration t.
Initialise the global iteration counter for each round do
Select Ny clients uniformly at random from (I.....N) without replacement. Let We C
{I...., N be the set of the participants (/u+l N2).
(Client update) For each of M iterations,
1. Perform an SGD update for the block uze That is
: where Tut Fet
(86)
Ttt
where ft is the minibatch version of f defined on the minibatch data batch, (so that Ebatch. (x)) = fx)).. Note that this update actually done independently over the participating clients E uu due to the separable objective, i.c., from (4) to (
2. < 1++
(Server update) For each of My iterations,
1.
Perform SGD update for the index (singleton block) 0. That is,
x1+1
[o0 [=+t!:='' where zo zo
1 Vof(('). 4
(87)
2. 1+1+1.
end
general non-convex globally, we can regard it as a convex function within a local neighborhood. as is usually assumed in non-convex analysis (Bertsekas) 2016) and other FL analysis (Li et al. 2020).
Assumption 2. For all t, fx(x) has Lipschitz continuous gradient with constant L.
Assumption 3. For all t, [+)ll S some constants.
Ry and ||> x'll S D for any x, x', where Ry and D are
Theorem D.1 (Convergence analysis). Let nh L + Vt, Ms M.N, and I" = T El x' by following our Fl. algorrth. With Assumptions 1-3, the following holds for any T:
N
- N N VT+I D2 + R2, 0
(88)
N₂ where x² is the (local) optimum, and the expectation is taken over randomness in minibatches and block selections {ur] )T-
T
Remark. Theorem D.I states that 전' converges the optimal point x in expectation at the rate of O(1/VB). This convergence rate asymptotically equals that of the conventional (holistic, non-block. coordinate) SGD algorithm.
To prove the theorem, we note that the algorithm Alg. [ overall repeats the following three steps per round: i) sample the subset of clients u, from {1...., N) with |ue] = Ny. ii) update z.. for M iterations, and iii) update zg for Ms iterations. Thus in the long-term view, we can see that the algorithm proceeds as follows: At each iteration I, we select u, as
(0)
with prob. Ms
M+Ms
14 =
(89)
Size-Ny subset uniformly from (1...., N} with prob. M+Ms
and update the iterate as
1+1 [ ; _ where 2. xt, f(x'))
(90)
T)t
We will use this long-term view strategy in our proof. Next we state the following lemma that is motivated from (Wang& & Banerjee 2014), which useful in our proof.
Lemma D.2. Assume 7)1 L, and Je is (locally) convex with Lipschitz continuous gradient with constant L. For any subset C {0, 1, 1,.,,,..NN. NN, we define x1+1 as:
1+1 x Vufe(x') and 2'1 x._0
(91)
ni
20

---

Under review as a conference paper at ICLR 2023
Then the following holds for any (Vufela'), x Iz) S Th (11x 2²11 || 21+2112)
R2
(92)
2
2(11: T.)
Proof of Lemma D.2 By definition, Vufz(z") + (e(( xx,,) = Then for any x; we have
(VuSela' ++1 Fn) = 7/ (1 F x-
x++ Fnn)
(93)
1t (1 - x1,113 || - x/1 1| 2 || ++ x 12)
(94)
Tt (1 x² ||
|| - 11 || |ج ++1 x
(95)
Note that (95) follows from 944 since x² and x only differ indices u. Since fo has Lipschitz continuous gradient,
I
Se(x/t1)- - Se(x") & (Vust(x") x - xxy) +
|
++1 ='11
(96)
2
I
(VuJi(x")) =5 xrtt! -xr.)=
|
='1l13 VV ( (1)), x.a)
(97)
2
+ I
(1|= 1²11 || ++
||
++ - x
-2,, 112
1 2zy11
=
2
2
V VuJi(("),+e
Zn);
(98)
where we plugged 955 into
97 We rearrange the terms in (98 as follows:
(V (Vafelz'), nu) S
1t (1 201-21112--- 1 11)
L -
|| ++1 x2,113
2
+ -fi(=t) - f((
(99)
Due to the convexity of fú, we can bound the last term in
99 as
f((x') - fo(2" (1x s (Vf((x'),
1
(100)
= (V.Je(v"), ) - 21t1)
(101)
$ 2a
/V + -1 x #11
(for any a 0).
(102)
2
Plugging 102) into 99) and choosing = 17t - -I(> 0) yields:
(uf((T')1
(V..i(x'). o x..) <
2 (|| 211 | 112)
(103)
2(1: L.)
Applying Assumption 3 of the bounded gradient norm completes the proof.
Now we are ready to prove our convergence theorem (Theorem D.I
Proof of Theorem D.7 We first aim to bound (Vf(x"), xx - z²., as it upper- of f(x') fx)
for convex f. Note that by conditioning on x², we can only deal with randomness in minibatch at (batch;), that is,
(Vx())) x' x) = Enatch, (VVe(('') x)]-
(104)
Further conditioning on batch, leads to:
(Vi((z'), - x) (Vof((''), =o)> + Σ (V.f(zz), xt xz).
(105)
Here we aim to rewrite the summation term in (05) terms of size Ny blocks. To this end, let us consider:
Σ (V.fi(''') xt xu).
(106)
uE2N.N,
where 2. is defined as the set of all subsets of [1,..., N} with size N, and no repeating elements. For instance, 25.3 contains {1,3,4 4) and (2, 4, 5), among others. Obviously |2, N
(,)) and
each particular index i {1....,N} appears exactly (N,=1) times in the 106). Thus we can
21

---

Under review as a conference paper at ICLR 2023
establish the following identity:
N - 1
Σ M
Vufe(z"), 24 In)
V at 22)).
(107)
N 1
We plug (17 into (OO)) and apply Lemmal D.2
(Vi( x' x) (Vole(x'), 20 -ro) + N
Σ
(VnJ(( xm)
(108)
(N
uE2N.N,
R2
|0 s'11 - -1| 1|--- 11 <
+ 2( L)
2
R2
Σ
(Z 21 - -11==1111-- -111--110------------1111 1 batch,) )1²) + 2(n -) (109)
(N
NE2N.N, where we define x² (u, batch.) for any subset u S. {0, I..... N) as: - == xt, - (1/m1e ) V (x")
and x''H == x² u. Although we can simply use x=+1, here we use this explicit notation to emphasise
dependency of g++1 on i and batch,, By letting g( xt, r(t1): == 2 (I| - x²1 - || - x 12)) + R7 2(m I)
we can express 109 succinetly yield:
(Vf(())) x)) S g(x,,', xt
++1 (0, butch:)) +
Σ NE2N.N,
g12, x', xttt (u, batche)).
(110)
For the second term, we use the uniform expectation to replace the sum (i.e., EuegMiNG v(u)=
(N)) B. [w]) for any function v) Using (N)) (x)) /x)(() N/NS, (Vf((a') xt - x) S g(x, ,
N E
x (0, batch,)) + N, Muz [g(x, x², + batche)),, (111)
and the right hand side of 1111 can be written as:
M + Ms
Ms
M
g(x, x", x. + batch)) +
M + Mg E N.N, BunzN.N, [o(a,, x at (u, batch)))
Ms
M + Ms
(112)
specification of Mg MMN2, Note that the expression inside the parentheses is exactly the expectation of g(x, x², -'+ batche)) over the random index set 44 following our client-server selection strategy in the long term view, that is, (89). Then we have the following result:
where
M + Ms
(V(((')) ) x0 x)
Ev, (g(x,, xf xx batch)))
(113)
Ms
where ue follows (89
As we have conditioned all quantities on batch,, now take the expectation over batchi. (V x) = Esatch, [VV fe(2"()), (x'), x)]
(114)
Ms M MM Ms Esatch..
R7 Te (1|= - 2' 22 || 111) + 2
(115)
2
where we drop the dependency in x ( batche) in notation, and use v)- = L + Vt. Since f is convex,
1(x²) - f(x) < (VS(x'), - x),,
(116)
and taking the expectation over x² leads to:
[[f(=')| f(x) <
M + Ms
Th mh (Ellc x' 111 - - (1 + 2V1
R2
(117)
Ms
2
By telescoping ( Et and using
M+ N+N, Nr
we have:
Ms
1
f(z*) -f())s N+ N₂
1
Σ
T
ne (B|zz 2'11 112 - x(+1112 + R2
Σ
1
N₂
2
20t
(118)
22

---

Under review as a conference paper at ICLR 2023
There are two sums in the right hand side of (118), and they can be bounded succinetly as follows. First, we use the simple calculus to bound the second sum: Σ 2/7 S J 2
+ S VT.
Next, we let a := B=|l 2 12 and the first sum is written as: 77 (a1 02) +... tpr (ar ar+1)
EL (ne TT-1)a; aa 7r ar+1 by letting no 0. Using o₂ S D from Assumption 3, this sum is bounded above by D2 (t= (n+ 7h-1) = D² (Mr - 7u) (vT + D)D2. Plugging these
bounds into 118 and applying Jensen's is inequality to the left hand side (i.e., E(
(L(((IT) x')] El (z")) yields:
N N N
VI+L 2
D2 + RRVT
B[s(=")) ((x)
(119)
T
for any 2, which completes the proof.
E GENERALISATION ERROR BOUND
In this section we will discuss generalisation performance of our proposed algorithm, answering the question of how well the Bayesian FL model trained empirical data performs on the unseen data points. We aim to provide the upper bound of the generalisation error averaged over the posterior distribution of the model parameters (ø, (0: (.)) by linking it the expected empirical error with some additional complexity terms.
To this end, we first consider the PAC-Bayes bounds (McAllester 1999 Langford & Caruana 2001
Seeger 2002 Maurer 2004), naturally because they have similar forms relating the two error terms (generalisation and empirical) expected over the posterior distribution via the KL divergence term between the posterior and the prior distributions. However, the original PAC-Bayes bounds have the square root of the KL in the bound, which deviates from the ELBO objective function that has the sum of the expected data loss and the KL term as it is (instead of the square root). However, there are some recent variants of PAC-Bayes bounds, specifically the PAC-Bayes-) bound, which removes the square root of the KL and suits better with the ELBO objective function (See TTiemann al. 2017) or Eq. (5) of (Rivasplata et al. 2019)).
To discuss it further, the objective function of our FL algorithm 3 can be viewed conven- tional variational inference ELBO objective with the prior p(3) and the posterior q(3), where 3 = {Ø, f1.... ....... ON' } indicates the set of all latent variables in our model. More specifically, the negative ELBO (function of the variational posterior distribution q) can be written as:
BL = [0x 1 (L(((())lp(3))
(120)
where i, (.3) is the empirical error/loss of the model / on the training data of size n. We then apply the PAC- MAC-Bayes-A bound (Thiemann et al., 2017 Rivasplata et al.) 2019); for any A E (0. 2), the following holds with probability at least 1--00 d:
Byce) [(3)] S 1- A/2 E Bots) (T. + X(1 (11-17) 2/2)
1
(I(((())lx(3)) - 1og(2
(121)
where 1(3) is the generalisation error/loss of the model 3. Thus, when À A= 1, the right hand side of [121 reduces to -2 ELBO(q) plus some complexity term, justifying why maximizing ELBO with respect to can be helpful for reducing the generalisation error. Although this argument may look partially sufficient, but strictly saying. the extra factor 2 in the ELBO (for the choice problematic, potentially making the bound trivial and less useful. Other choice of > fails to recover the original ELBO with slightly deviated coefficients for the expected loss and the KL.
1) may be
In what follows, we state our new generalisation error bound for our FL algorithm, which does not rely on the PAC-Bayes but the recent regression analysis technique for variational Bayes (Pati et al. 2018 Bai et al. 2020). It was also adopted in the analysis of some personalised FL algorithm (Zhang ct al. 2022) recently.
E.
GENERALISATION ERROR BOUND VIA REGRESSION ANALYSIS TECHNIQUE
We begin with the regression- regression-based data modeling perspective and related assumptions/notations.. We denote by (z. y) the true data distribution for client i (i = I....,N.. ). We assume that the target is real vector-valued (y E IR Sv). and there exists a true regression function I SR - RRS. for each
23

---

Under review as a conference paper at ICLR 2023
1. That is,
(yyy)) - ["); (2))
(122)
where o² is constant Gaussian output noise variance. Let D; (xi, ~ pz(x,)) v) be the Li.d. training data size n for each client i.
In our FL model, we assume that our backbone network is an MLP with L hidden layers of width M, and all activation functions o(-) are Lipschitz continuous with constant 1. The parameters of the MLP are also assumed to be bounded, more formally, the parameter space @ is defined as:
0 E Q {0 E R ||0||=0 I B, MLP with L layers of width M).
(123)
Note that G = dim (f) and B is the maximal norm bound. The MLP defines regression function fo R. + R and the i-induced predictive distribution is denoted as:
(o(y)z) .N Ny; fo(z), o?1),
(124)
where we assume that the true noise variance is known.
For notational convenience, we denote by f(X') the concatenated vector of f(x) for all E X i.e., f(x') = [_(=)]=xx', where ((-) is either the true (')- or the model fo(-). Extending this notation, simply writing f or fo means
(population) responses, that is, f [f()))EERS, and fe similarly. For instance, llfe fo stands for the worst-case difference, namely max,ERs= [Ifo(x) - fl))]l. As a generalisation error measure, we consider the expected squared Hellinger distance between the true P and the model Po. Formally,
(Po, P') EE-P(e) ( (Pooll))))) P' (E 1- exp (- [1fo(z)) fx()
80?
(125)
More specifically, we will bound the posterior-averaged distance N Σi=1 (0) Id (Po., Pi)], where
{or )/= is an optimal solution of our FL-ELBO optimisation problem (We showed in Sec. D
that our coordccooe FL algorithm converges to this optimal solution in O(1 rate).
Theorem E.1 (Generalisation error analysis). Assume that the variational density family for (0,))
rich enough to subsume Gaussian. The optimal solution (1 (0.)) (x)) of our FL.EEIBO optimisation problem
3 satisfies (with high probability):
(
E Co (Po,s Peo,PP') P')1 o ((
N
+ + +
C²
I' N
ΣΣ
(126)
<=1
where C, C' 0 are constant,
X mingge 11.5o 113 is the best error within our backbone O,
In G(L log M + log S. (5 n
G
(127)
n
and 6, VTn log (n) for 8 constant.
Remark. Theorem implies that the optimal solution for our FL-ELBO optimisation problem (attainable by our coorrdiatee FL algorithm) is asymptotically optimal, since the right hand side of (126) converges to 0 as the training data size + co. This is easy to verify: as - co, In + obviously, accordingly 6n - 0, and the last term v Σ, Ar can be made arbitrarily close to 0 by increasing the backbone capacity (MLPs as universal function approximators). But practically for fixed as enlarging the backbone capacity (i.c., large G. L. and M) also increases en and rn. it important to choose the backbone network architecture properly. Note also that our assumption on the variational density family for gi (0,) is easily met; for instance, the families of the mixtures of Gaussians adopted in NIW (Sec. 3.1 and mixture models (Sec. 3.2) obviously subsume a single Gaussian family.
Proof of Theorem E.1 We first aim to link the variational ELBO objective function to the Hellinger distance via Donsker-Varadhan 's (DV) theorem (Boucheron et al.) 2013), motivated from (Bai et al. 2020 Zhang et al. (022). The DV theorem allows us to express the expectation of any exponential function variationally using the KL divergence. More specifically, the following holds for any distributions q and any (bounded) function h(()):
log Bp(ss) (e ((2)) max (Eycs) kL(g))))
(128)
q
Note that the optimal posterior on & (i.c. 9" (x)) does not appear here since it only affects dP Pt) implicitly. See our detailed analysis/proof provided below.
24

---

Under review as a conference paper at ICLR 2023
Here set p(z) = p(0)1)), 9(2) == (:())). h(2) == log (8s))) where
(1())) = exp (Po, (D:), (D())) nd (Po;, Pi)) and
(129)
Pe: (P.)
Ln (Po, (D.), P' (D())) =olog (D,)))
(130)
and we have the following inequality that holds for any Ø:
log (p(0,10) (0.)] 2 (0:) ny (0:)] - By taking the expectation with respect to (()) and rearranging terms, we have n (e, (d? (Po((PP) P'] & Ea.c, (Po (D.), (D())) + Ey(a) (L(((0))))(0)))))) Ey(o)] log (p [0)]].
(131)
+
(132)
For the last term of the right hand side, we use the bound (x(e)(((0) S. eC from the regression theorem (Pati et al. 2018), which holds for any distribution s(0) with high probability. The details and proof of this bound can be found in the proof of Theorem 3.1 in (Pati al. 2018). Applying this bound and telescoping over = -1..... N yields:
N
N
Bmto (Po, P')] S
E (u
+
N
Σ
Ey(s) [KL( + NCne,
(133)
We add (L((())lp(()) the right hand side, which retains the inequality since KL divergence is nonnegative. Then we have the following result that holds for any with high probability,
N
n
E (e (o (Poz) P')] & In(a(o), (.)00,)) + Cne
(134)
N
where In (0)), {g. (0:) )==)) equals:
E LP (D.), Ey(s) [L((((()))
(L(()))
(135)
which exactly coincides with our FL-ELBO objective 3 up to constant, by the factor of 1/N.
Next, we define (0:) and g(Ø) as follows:
G A,
ã (0.) (N(0;;; 07, o with 0, arg min IIʃo 1'11.
o
where
(136)
8EB
8n
A = -ogg(3,, (2BM))(()))))) ((5, + + BM -1
2
+
(2BM)2- -1
(2BM 1)2
((6) = arg min
Σ
B KL (L [ + (L((())lp(o)))
(137)
4(
Since (l (0.)) (o)) is the optimal solution of the FL-ELBO optimisation problem, it is obvious that (o), (a) (0.)), L,, (ō()). fē, (0,)) if the variational density family for q. (0,)) is rich enough to subsume Gaussian. Now we look closely at T., (ā()), (0.)0) ) and we note that the last two terms as per 135) are constant (i.c., not a function of data size n1). That is,
N
Ei(s) KL -
(138)
N
for constant Ò. Then we write
N Ea (- (D),, + č.
L, (ū(0), (Q.)))
(139)
N
We bound the expected -,, term in (19 using use Lemma E.2 below which states that with high probability,
( (D., P(()) < C'n(r, + +
(140)
Althouhh I Lemma R.2 can be found in the proof of Lemma 4.1 of (Bai at. 2020 we state this lemma
clearly with separate proof for self-containment.
25

---

Under review as a conference paper at ICLR 2023
for some constant C' > 0. Plugging this bound into (139) 1. we have the following derivation where
we start from (134) with ( = (Ø) and g:(0.) = ((.))
N
72
E (o) [d (Pe., z
+ Oner
(141)
N
s Ln(ā(o), (1.0))))) + Oner
(142)
1
S S + C'n
na
Σ x1 N
+ Cne
(143)
By dividing both sides by we complete the proof.
Lemma E.2 (From the proof of Lemma 4.1 in (Bai et al. 2020)). For (.().) defined
12 as in 136
and ras X defined as in Theorem E. the inequality (400) holds with high probability.
Proof of Lemma E.2 From our regression model assumption 122 and 124
(q,(0.)
(PPo Ln
((.))))))))) = Eq,0, log (DD) - og Po. (D.]
(144)
1
fo. (X')
E E f(
(145)
20?
1
s'x +2 Ea, (Xx)),((X" So, (x9))
(146)
2o?
AA
R
We first work on R1. Since
[x+)) - -xa) fxx")]
< n [|50, 1' S 2n 2n(11o, -00 + +1150 - 8 11122
(147)
we have
R 2n (0, 11. + 2n 11fo; 1/
(148)
< 2n(rn + )1).
(149)
where in )49), we use the definition of X; and the fact (., - soz 123 Snn from Appendix G
in (Chériel-Abdellatii 2020).
Next, we bound Rg. Since (Y* f(x')) ~ N'00, 1) and independent of 0,, we can let e :=
Y - f(X)) for N(0,0 o. Then
Rs cT "Eg.o.
,N(0, NNO0
(150)
where of = [B [ So. (x*) 112 Applying Jensen's inequality on the convexity of || 11.
Cj S lf (X+) fe fe, fx (x = R1. Due to the property of Gaussian, there exists some constant
C6 such that R2 S C. cy S C. R1 with high probability. Plugging these bounds on R₁ and R₂
back to (16) leads to:
1
1 20% R
BR(())) [- 1n (Po, (D.), PD.))]
=
20
(R1 + +22))
20
(151)
+ 20%
< o7
n(rn + Ai).
(152)
Letting C²
1+22/
(constant) completes the proof.
F ADDITIONAL EXPERIMENTAL RESULTS
F.1 MORE RESULTS ON CIFAR-100 & CIFAR-C-100
We test our mixture model with different mixture orders K = 2, 5, 10 on CIFAR-100 (Table and CIFAR-C-100 (Table 4.. In the last columns of the tables, we also report the performance of the centralised (non-FL) training, in which the batch sampling follows the corresponding FL settings.
12 In Bai et al. 2020). they defined q. (0.) as a spike-and-slab model to deal with sparsity. Essentially it is a mixture of two components, selecting (0,; 0,, of when 0; entries are non-zero and selecting the delta function at when 8; entries equal zero. Without loss of generality (or under mild numerical approximation) we assume all entries of 0e non-zero, which makes (0,) Gaussian equal to ((0,; 0 2 in 136
26

---

Under review as a conference paper at ICLR 2023
Table 3: (CIFAR-100) Global prediction and personalisation accuracy. Mixture order K varied. Comparison with centralised (non-FL) training.
(a) Global prediction performance (initial accuracy)
FL settings
Mixture (Ours)
NIW (Ours)
Fed-BABU
Centralised
5
K
K 5
10
1
19.76=0.11
19.3720.30
16.9720.13
48.35
12.35
52.21 10.19
0.1
10
29.02
29.02
29.08
29.90
27.9 27.93
36.87
100
1
57.80=0.10
52.9420.36
52.24
51 51.92400.0
18.17 18.17200.56
58.5020.52
1.0
10
29.53 #t.422 40.42
30.55 30.55 10.15
30.85 30.85 10.13
30.24 400.49
28.67 28..7
48. 18.16
1
37.54 20.25
38.07
39.96
39.97-0.58
35.04
28.55
0.1
18.99 10
#0.13
18.93
18.95
18.54
9.92
10
50.40-2.111
49.52
49.82
49.51 1:0.36
45.41 1:.11
78.37 1.088
1.0
10
22.87=0.41
23.59
23.69
24.28:11.04
21.92
4.40
(b) Personalisation performance
FL settings
Mixture (Ours)
NIW (Ours)
Fed-BABU
Centralised
f
K =
K
= 10
1
54.16-0.500
56.17-0.16
54.93
55.8310.47
50.43-0.99
53.18 53.18
0.1
10
36.68#0.37
36.32=0.27
37.30+0.64
37.34+0.38
35.15+0034
38.20
100
60.36*0.88
58.8220:37
58.16 58.1640.26
58.3240034
55.87400.91
58.49440.50
1.0
10
35.92*0.17
36. 36.22#0.17
36.44
35.91
35.58++0.24
48. 18.17
79.41 #0.24
79.19
49.19
77.44
4..44
42. 42.04 12.04:41.35
0.1
10
67.35 1.002
67 67-57#0.62
67 #0.40
67.33:20.26
66.24 40.63
17.40
10
1
82.71
#0.37
81.03 0.33
79.91 4:0.5
80.92:0.119
78.92
78.43:10.90
1.0
10
67.78
66.74
60.50
67.30
66.25##0.46
5.13
Global prediction (s= 100)
Personalisation (s 100)
50
40
Accuracy
Accuracy
30
Ad-BABU
Fed-BABU
20
F
Preset (Baseline)
Preset (Baseline)
T
Ensemble (Baseline)
20
Ensemble (Baseline)
10
Mixture (Ours)
Mixture (Ours)
10
10
K (the number of networks)
K (the number of networks)
Figure 3: Comparison between our mixture model and ensemble baselines (K varied) on CIFAR-100.
That is, at each round, the minibatches for SGD (for conventional cross-entrope loss minimisation) sampled from the data of the participating clients. The centralised training sometimes outperforms the best FL. algorithms (our models), but can fail completely especially when data heterogeneity is high (small s) and is large. This may be due overtraining biased client data for relatively few rounds. Our FL models perform well consistently and stably being comparable to centralised training on its ideal settings (small and/or large s).
F.2 COMPARISON WITH SIMPLE ENSEMBLE BASELINES.
Our mixture model maintains backbone networks (specifically {r,))= where the mixture order K is usually small but greater than (e.g., K = 2). Thus it requires extra computational resources than other methods (including our NIW model) that only deal with a single backbone. As a baseline comparison, aim to come up with some simple extension of Fed-Avg (McMahan et al.) 2017) that incorporates multiple (the same K) networks. Here are the detailed descriptions of the baseline ensemble extensions:
27

---

Under review as a conference paper at ICLR 2023
Table 4: (CIFAR-C- 100) Global prediction and personalisation accuracy. Mixture order K varied. Comparison with centralised (non-FL) training.
(a) Global prediction (initial accuracy) on test splits for the 10 training corruption types
FL settings
Mixture (Ours)
NIW (Ours)
Fed-BABU
Centralised
5
1
81 $1.2240.14
80.312144
81.30=0.21
81.2320:79
79.45220.71
$7.68-0.48
0.1
4
67.69 67.69
65.81
65.81=0.84 00.84
61.37=0099
66.72#1.46
63.58
76.82
100
1
91.2 91.26:10.3
86.8120.22
86.41 46.55
87. 87.0420.33
86.84
90.22-0.57
1.0
4
73.58 :..20
74.55 20.22
74..33#0.81
75.03 40.45
71.03
$7.71=9.90
1
78.63
79.36-10.2
78.23440.3
78.45 78.45+40.00 1.040
77.44
84.97=0.22
0.1
65.08=1.75
63.52#0.45
63.08=0.55
63. 13.13300.21
62.65+0.12
70.76=2.83
50
89.31 1:0.17
28.24
87.81 10.65
87.02
86.44 67.66
89.92-0.47 0.47
1.0
4
70.33
70.19
71.87=2.31
74.18
86.40=1.00
(b) Global prediction (initial accuracy) on the original (uncorrupted) CIFAR-100 training sets
FL settings
Mixture (Ours)
(Ours)
Fed-BABU
Centralised
5
=
K
K = 10
1
41 41.55-10.11
36.99
37.75 20.50
37.7 37.78
34.76
35.10=0.86
0.1
30.84+0.07
30.60+0.27
30.51 #0
29.13
28.31 +0.28
33.6/=1.08
100
4..3220..2
38.35 0..73
20,00.00 96
38. 1.022
35.58:00.11
92.78=0.22
1.0
4
30.67
30.40#00.44
30.57#0.25
30.95 30.95+00.28 +0.28
28.60
28.95=2.00
1
11..04 40.14
40.14
36.41 #0.47
36.03 0..22
38.3 38.32 00007
35.44 4.055
35.71 =0.13
0.1
4
40.36
31.30=0.34 00.34
31.19 #0.26
31. 31.06= +0.58
29.68 29.68=20.08 +0.08
35.28=1.47
50
1
41.64 64 40.21
20.44
39.19 0..37
38.93 1.115
36.09 2.228
33.49 0.11
4
32.17
30.68#0.46
31.29
32.466#0.20
29.28
29.50=1.24
(c) Personalisation performance on the 9 held-out corruption types
FL settings
Mixture (Ours)
NIW (Ours)
Fed-BABU
Centralised
K = 2
K =
K = 10
72.63
74.1643.04
74.43
75.65
69.93#1.24
88.76
0.1
62.74
61. 1.65
60.16
64.41=2.66
60.33
80.47
100
83.62
84.88:10..85
83.01 1:779
82.58 0..91
77.55 42.05
87.55 42.93
1.0
6..8
67.35
64.34
15.12
53.25#2219
76.58 12.63
75.50 40.74
67.33 2.33 4.283
4..32
65.34 63.14
$3319
59.47 42448
84.08
84.08+11003 1.03
0.1
4
41.20
46.39
45.95
45.43
44.74 1.30
74.14
50
1
81.46
81.77 43.11
78.03 78.03+1.05
74.62
67.43 42.58
87.57
1.0
50.84
40.50
45.07
48.78
40.87#3.01
81.46#1.11
1. The server maintains K networks (denoted by 0! 8!....,KK)
2. We partition the clients into K groups with equal proportions. We will assign each 0* to each group j (j = K).
3. At each round, each participating client i receives the current model (i(+) from the server, where j(i) means the group index to which client i belongs.
4. The clients perform local updates as usual by warm-start with the received models, and send the updated models back to the server.
5. The server collects updated local models from the clients, and takes the average within each group j to update 65
6. After training, we have trained K networks. At time, we can use these K networks in two different ways/options: (Preset option) Each client uses the network assigned to its group. i.c., gz(1), for both prediction and prruuningpersonnaisationnn, (Ensemble option) We use all K networks (as ensemble) for prediction and finetuning.
28

---

Under review as a conference paper at ICLR 2023
Table 5: Running times (in seconds) on CIFAR-100 with (s 100, f
1.0,
1) setting.
Fed-Bin BABU
NIW
Mix. (K 2)
Mix. (K 5)
Client
0.283
0.362
0.481
0.595
Train
Server
0.027
0.267
1.283
1.434
Global prediction
0.021
0.025
0.040
0.060
Porsonalisation
1.581
2.158
2.421
2.766
Note that K = 1 exactly reduces to Fed-Avg (or Fed-BABU). In Fig. 3 we visualise the performance of these ensemble baselines, compared with our mixture model for different K 2, 5, on CIFAR- 100 with (ʃ 0.1, 1) setting. It clearly shows that these simple ensemble strategies are prone to overfit. The result signifies the importance of the sophisticated negative log - ses-sum-expp regularisation in the client/server updates in (2 and
23
in mixture model.
F.3 RUNNING TIMES
Although our models achieve significant improvement in prediction accuracy, we have extra com- putational overhead compared simpler FL methods like Fed-BABU. To see if this extra cost is allowable, we measure/compare wall clock times in Table 5 where all methods are tested on the same machine, Xeon 2.20GHz CPU with single RTX 2080 Ti GPU. For NTW, the extra cost in the local client update and personalisation (training) originates from the penalty term in (17,, while model weight squaring to compute Vo in (18) incurs additional cost in server update. For Mixture, the increased time in training mainly due to the overhead of computing distances from the K server models in (22) and (23). However, overall the extra costs are not prohibitively large, rendering our methods sufficiently practical.
G (REVISION) ALGORITHNS/PSEEDDCODDS AND COMPUTATIONAL COMPLEXITY ANALYSIS
ALGORITIMS//SSUDDOCDDD
We provide pseudocodes for our algorithms:
1. Training algorithms
General framework Alg.
Hormal-Invers--Wihhee case in Alg.
Mixture case in Alg. 4
2. Global prediction algorithms
General framework in Alg.
Normal-Inverse- Wishart case in Alg.
Mixture case in Alg.
3.
Personalisation algorithms
General framework Alg.
Normal-Inverse- Wishart case in Alg. 9
Mixture case in Alg. 10
G.2 COMPUTATIONAL COMPLEXITY ANALYSIS
easily analyse the computational complexity of the proposed algorithms. They are summarised in Table (training complexity with communication costs), Table (global prediction complexity), and Table (personalisation complexity).
Based the detailed algorithms in Sec. G.I
29

---

Under review as a conference paper at ICLR 2023
Algorithm 2 Training algorithm: General framework.
Input: Initial parameters Lo the variational posterior g(:; Lo).
Output: Trained parameters T.p- For each round r = 1,2..... R do:
.,
1. Sample subset I of participating clients ((Z) Ny N).
2. Server sends Lo to all clients I.
For each client
in parallel do:
Solve (by SGD) with Lu fixed:
min (0; log (D]0:)] + Begtss))) [L: L.) [=)))(((()))]]
Initial L, can be either copied from La or the last iterate if the client able to save L, locally.
4. Each client i E I sends the updated L, back to the server.
5. Upon receiving {L, G bez, the server updates Lu by solving (with tez fixed):
min KL(g(0; Lu)llp(o)) Lo
N
Σ
NI
IEZ
Algorithm 3 Training algorithm: Normal-Inverse- Wishart case.
Input: Initial Lo (mo. Vo) in ((;; La) NZW((; (ma, Vo, = |D] + 1, = [D] d 2)) where
|D|= Σ i=1 [D_] and d = the number of parameters in the backbone network p(ylz, 0).
Output: Trained parameters Lg (mo. Vo).
For each round r ..,. do:
1. Sample subset participating clients ( = N, s N).
2. Server sends L-g = (mo, Vo) to all clients i E I.
3. For each client I in parallel do:
Solve (by SGD) with Lo (mo. Vo) fixed:
min - log p( + _ (no + d 1)(m: - mo) T V/- 1m, mo), my
where m, is the dropout version (with probability 1 p) of m,,
Initial m. can be either copied from ma or the last iterate if the client able to save m, locally.
4. Each client E sends the updated Li m, back to the server. Upon receiving {m, Jzz, the server updates Lo = (mo, Vo) by:
5.
mi
P
N
Σ m4: VoL N+ (1+ Ne + mm(m))
N
Σ
p(mm m..p)
N+IN,
N₁ IEI
where p(mm,, m. p) = pm pmmmm
pm
mome
Algorithm 4 Training algorithm: Mixture case.
Input: Initial Lo = fr, ト in g(0; Lo) = II, Niyi e 2I) and / in the gating network g(2; ())
Output: Trained parameters T.g {r,}5= and B.
For each round ..2...., do:
Sample subset participating clients (II) N, S N).
2. Server sends T.p {r, ) and to all clients E I.
3. For each client I in parallel do:
Solve (by SGD) with Lo
{r> fixed:
K
||me - ryy
min E. (0.mm))) log p(()).) log Σ exp
where (;(0;; (000; e ef).
20²
Initial m. can be either the center of (ry) the last iterate if the client is able to save mn, locally.
8, SGD update in g(a;; (3) with data ((x,j")) where j" arg min; [|m, - ry 4, Each client E sends the updated L; = m. and 8i back to the server.
5. Upon receiving {mi and (B, the server updates To = {ri}5=1 by the one-step EM:
Ex mm-rr 1//(2035' (M-step) rz Ny Eier c(ili) m.
e -11m--r511//007
(E-step) k(jlk)
NV N. EIEz e(ili)
and updates 3 by aggregation: 8" Nr Eier Br.
30

---

Under review as a conference paper at ICLR 2023
Algorithm 5 Global prediction: General framework.
Input:
Test input x². Learned model Lo the variational posterior q(:; Lo).
Output:
Predictive distribution p(y"|=*' DI:N).
Sample gfe) p(0))) q(d; La) de for S.
2. Return p(y" [2" D1.x) Σ E ply" 1x",, ,
Algorithm 6 Global prediction: Normall-nverre- Wishart case.
Input:
Test input x² Learned model Lg (mo, Vo) in g(0; Lo) NIW((d; (mo, Vo, Vo,lo.no (o. Output: Predictive distribution p'y" |x ,,.N).
Sample g(+)) ~ Loo-d+1 00 mn, In("0 (1190-4+1)
(1++1iV6 for = for.:,,...,55
2. Return p(y"ee l0,, DI.N) R E. p(y" 1x" )
Algorithm 7 Global prediction: Mixture case.
Input:
Test input x²" Learned model Lo {rybf= in ( To) W1, and / in g(n; (8).
Output: Predictive distribution p(y"(x', D.:N). Return p(y" lx",, D1.N) Σド= (x(x") ((y"""",r)))
Algorithm Personalisation: General framework.
Input:
Personal training data D Test input x². Learned model Lg in the variational posterior g(0; Lo).
Output: Predictive distribution p(y""="" [="", D" D1.N).
1. Estimate the variational density v(0) ze p(0 D e²) by solving (via SGD):
min Ee[ss( logp( (0)] - (I(()())]p((()) )), where o arg max (d;; To).
2. Sample (6*) c(0) for S.
3. Return p(y" wo" D", D1.N) = 3 E. p0 ((0)).
Algorithm 9 Personalisation: Normal- Lanemaal-Lneers--WWieaa Wishart case.
Input: Personal training data D Test input x".
Learned model Lo (mo. Vo) in g(tb; Lo) = NIW(( (mo. Vo, lo no)). Output: Predictive distribution p(y"|x* |="" D" D1.N).
Estimate m v0; m) = II: P (([1]] m(1] 22) (1- -V(00))0 22)) by solving (via SGD);
min log p( (m) L (no d 1) mo) Voa (m mo),
where in is the dropout version (with probability - p) of m. 2. Return p(y" le,, D", Di.N) = p(y"1=" m).
Algorithm 10 Personalisation: Mixture case.
Input: Personal training data D Test input x".
Learned model Lo {r, in g(0; Lo) 11, Nu,ii,,, Predictive distribution p(y"|>" !=." D", DI.N).
Output:
1. Estimate m in v(0; m) = (6; m, (1 by solving (via SGD):
m rryl
min E log.ll - log
exp
20
1=1
2. Return p(yPlet We DP, D = ply pam, m))
H (REVISION): ADDITIONAL EXPERIMENTS
H.1 CIFAR-100 WITH THE EXTRA FOUR FL METHODS
The results are summarized in Table where we use the following hyperparameters for the competing
methods. FedPA (Al-Shedivar et al. 2021): shrinkage parameter 0.01; FedBE (Chen & Chao
31

---

Under review as a conference paper at ICLR 2023
Table 6: Training complexity of the proposed algorithms (NIW and Mixture) and Fed-Avg. All quantities are per-round, per-batch, and per-client costs. In the entries, d = the number of parameters in the backbone network, F time for feed weed-forwrrd pass, time for backprop. and N, the number of participating clients per round.
Communication cost
Client Local update
Server update
Server Client
Client + Server
2d
d
F' + B + O(d) (O(d) from quadratic penalty)
NIW
O(Ny .d)
(sent: ma, Vo)
(sent: m,)
Mixture
(K 4)d
2d
2(F B) O(K d)
O(K - dd)
(Order K)
(sent: {r, {r;))=1,8) B)
(sent: ms, (.)
(O(K dd from log-sum-exp)
d (sent: 0) (sent: 0:)
d
O(N, dd) (aggregation)
Fed-Avg
F+B
Table 7: Global prediction complexity of the proposed algorithms (NIW and Mixture) and Fed-Avg. All quantities are per-test-batch costs. In the entries, the number of parameters in the backbone network, F time for feed- nerd-forward pass, and the number of samples ((*) from the Student-t distribution in the NIW case (we use S = 1).
Per Per-test-batee complexity
F + O(S .d)
NIW (0(S d) from the cost of (-sampling)
Mixture OOrder K)
(K + 1)F
(a forward pass for the gating network)
Fed-Avg
F
Table Personalisation complexity of the proposed algorithms (NIW and Mixture) and Fed-Avg. All quantities per-train/test-batee costs. In the entries, d = the number of parameters in the backbone network, F time for feed-forward pass, and B time for backprop.
Training (personalisation) complexity
Test complexity
F+ ((d)
NIW
F
(O(d) from quadratic pemalty)
Mixture (Order K)
F B + O(K
(O(K- d) from log- sum-ep)xp
Fed-Avg
F +
2021):
the number of ensemble components 3; FedEM (Marfog et al. 2021): the number of base models 3.
H.2 MNIST AND DASHION-MNISS
The FL setting is follows: the number of clients N 100, the number of shards per client 5. the fraction of participating clients per round f = 1, and the number of local training epochs per (total number of rounds 20). FedBE and FedEM
round (total number of rounds 100) three component models. The backbone is an MLP with single hidden layer with 256 units. The results are summarized in Table I0 (MNIST) and Table II (Fashion-MNIST).
H.3 EMNIST
The FL setting is follows: the number of clients N 200, the fraction of participating clients per round f = 0.2, and the number of local training epochs per round 7 = (total number of rounds 300). We follow the standard Dirichlet-based client data splitting. FedBE and FedEM use three
32

---

Under review as a conference paper at ICLR 2023
component models. The backbone is standard ConvNet with two hidden layers. The results are summarized in Table
I (REVISION): EXTENDED RELATED WORK
General FL approaches. Perhaps the seminal pioneering work FL is attributed to FedAvg (McMa- han et al. 2017). which proposed fairly intuitive local training and global aggregation strategies with minimal training and communication complexity. A potential issue of divergence of global and local models due to the separated steps of local training and aggregation was addressed by model regularisation in the follow-up works (Li et al., 2018 Acar et al.) 202T), which is shown to help the global model converge more reliably. Recent approaches aimed to incorporate benefits from existing machine learning approaches including domain adapttationgecnrralistttoon clustering. multi-task learning, transfer learning, and meta- meta-learning. To deal with heterogeneous client data distributions, those works in (Peterson et al., 2019 /Zhng et al.) [02T] seapptaaioonGeneraliinnnn
|| [Sun et al. (02T) attempted to tackle the FL
Another interesting line of works aims to cluster clients with similar data distributions together (Briggs et al. 2020 Mansour et al. 2020 Along the line, the shared representations among the related simila lar clients can be modeled motivated from general multi-task learning (Smith et al. [017]] Dinh et al. (202I) I) Motivated
problem in the perspective of (multi- Domlain
from transfer learning, reasonable attempts are made to exploit the Idea of sauminn/transferring knowledge from related clients (Chen et al. 2020 Yang et al. 2020. Dinh et al. 2020 Li et al. 2021]
Last but not least, there have been attempts to the personalised FL methods based on meta learning
since the fientuning from the global trained model can be seen as adaptation to new data (Chen et al. 2018 Fallah al. 2020).
Comparison to existing Bayesian FL approaches. Some recent studies tried to address the FL problem using Bayesian methods. As we mentioned earlier, the key difference is that these methods do not introduce Bayesian hierarchy, and ultimately treat network weights @ as random variable shared across all clients, while our approach assigns individual 0, to each client i governed by a common prior ((). The non-hierarchical approaches must all resort to ad hoc heuristics or strong assumptions in parts of their algorithm. More specifically, FedPA (Posterior Averaging) (AI-Shedivat
et al. 2021) aims to establish the decomposition, p(0) DiNN) ) x II; p(0) Di) also known as product of experis, to allow client-wise inference/optimisation of p(0) D:). Unfortunately this decomposition does not hold in general unless we make a strong assumption of uninformative prior p(0) xx 1 as they did. FedBE (Bayesian Ensemble) (Chen Chao (02T) aims to build the global posterior distribution p(0) D1.v) from the individual posteriors p p(0 D.) in either of two ad-hoc ways: SWAG (Maddox et al. 2019)-like model averaging over clients, or a convex combination of the modes of the local method to approximate D(,N) from individual posteriors p(0) To combine the individual posteriors. they introduce the so-called global distribution w()), which essentially serves a regulariser that aims to enforce local posteriors p.) not to deviate from it, i.e., p.)) RS w(()) for all i. The introduction of w(() and its update strategy appears to be hybrid treatment rather than solely Bayesian perspective. Finally, FedEM (Marfoq et al. 202I) forms a seemingly reasonable hypothesis that local client data distributions base distributions (with different mixing proportions). However, although they have probabilistic modeling, mixture estimation, and base distribution learning under this hypothesis, this method not a Bayesian approach.
posteriors. PFedBayes (Zhang et al. (002)) can be seen as an implicit regularisation-bat
be identified as mixtures of fixed number of
GedGP (Achituve et al. bb2 aims extend the GP-Tree algorithm (Achituve et al. 2021a) to the FL setting via the shared deep kernel learning. To this end, the clients perform GP-T Tree kernel learning locally on its own data while the server aggregation simply follows the FedAvg algorithm to learn a global kernel. In this sense, the overall approach is quite different from our hicrarchical Bayesian treatment. FedPop (Kotelevskii et al. 2022): It has similar hierarchical Bayesian model structure as ours. But they split the backbone network parameters into those of the feature extractor (denoted by ø in the paper) and the linear classification head (3). In their model. the feature extractor weights & are shared across the clients (called fixed effects), and the client-wist wise classification head parameters 2 are sampled from B, i.e., 2' ~ p(2]8). Thus the client data D, is generated by 6 and 2². The main differences from our approach are in four folds: 1) The higher-level variables 3 and local variables 2f sampled from 3 are both restricted to the linear classification head part of the network, which makes imposing uncertainty in model parameters quite limited; 2) Moreover, they do not
33

---

Under review as a conference paper at ICLR 2023
FedPop
43.00 33.09 03.03.00.30
1. 02
48.5 18.36 1.302-044
28.09=0-17
35.01 35.0140.30 4.0.6
18.68 #0.12
44. 44.98
22.0620..0 40
FedPop
51 61.2220:37
35.69 35.09 40.47
55.85
35.36
74.97
66.70 +0.53
78. 78.96
06 00-5722044
28 28.31 03.31-m0
22.
28.
FedEM
43 33.8341.00
28 08.62 20.20.220
50 20-72
29.5120.12 12
37 37.28 20.26
18. 18.45##0111
47.52
20006
22.51
FedEM
47.28 68
41.11
54 4..65
35 35.21 12
75. 75.5240.50
67.11
80.65 32
06.34#012
29. 29.51
FedBE
42. 12.49-00.0 89
28 28.39=0.00
48. 08.6920-70
28. 28.60#1115
35. 17.1720..0
18.67
45.56 52
21.7320.0 61
FedBE
50.5721.0
35.51 14.1400000 62
56. 20.25200-2
34.92
76. 76.199
06.64 +0.25
79.8220.52
06 06.0620.54
79
(CIFAR 100) Global prediction and personalisation accuracy.
FedPA
42. 42.15
2.0.0
47.90 47.9620.11
28 28.89#0038
35 35.51 400.5
18. 18.59
45.08 45.08
#0.20
22 22.60
FedPA
49. 88:88 40
30
35. 06.06:0.000
55. 23-3.33
35. 35.21 #00.00 50
75. 76.76000.00
65. 87 10.21
78. 88.8420.116 16
23.23300.2 20
28. 28.05
66.
Global prediction performance (initial accuracy)
pFedBayes
37.23 88
28.21 11.11.00 42
448922.32
28.2520.81
$1.4940.18
17.93 68
39. 39.95
31.8520.30 50
pFedBayes
45. 13
45.34
53.0 #0.00
35.57#1002 02
+0.33
78. 78.08#0.26
66. 11.22
73.03
65. 65.82*h
28.
21
(b) Personalisation performance
Fed- Prox
41 $1.4940.73 75
27.11#0.11 10
47.6 47.6621 40
17.4320.5 38
34.73 24.734022
16.84
44 44.3320.31
19.5420.38 38
Fed-Prox
49.91 =0.78
1.0.22
55. =0000
33. 33.70=0-12
#0007
75. 10.00=0.00
64. 16 +0.73
78. 18.56=0.55
03.8 #00.1
$3.92-000
33. 33.92*
27.
Fod-Avg
40 01.8720.00 62
28. 25.26:20.19
47. 47.4420.2
28. 05.79200.0 68
48-48.0000 86
14.69 40.04 46
37 37.10=0.44
17. 17.3820.. 32
Fed-Avg
46. 82
33.57.20.00
53. 20.15
33 03.82#1001
70. #0.30-1.00
61.39 +0.27
76.98 68
63.81 20 40
27 17.48
Fed-BABU
42. 42.35-2..4 42
27.93 77.9340.2 28
48. 18.1720.66
28. 23.6740..51
35. 100440.000 56
18 54 40.377
45. 45.41
21 31.9220.00
60
Fed-BABU
50 43.0..03
55, #7.99
35. 35.58=0.22
75. 14.4 30
66.24*h #0.53
78. 02.2220.22
23
06. 20
35. 15.45-000
Table 9:
Mix. (K = 2)
49 00.3770.30
00.29
52. +2.94 10.36
30 00-55
38 05.0720.40
18 .5.95* 46.9540..13
49. 50.882
40.47
2)
56.17 17.11
36.32*h 36.3244..27
58. 8220.37
36 20.17
79. 70.
67.57 40.62
81. $1.03=0.33
66 76.7420.22
Our Methods
Our Methods
Mix.
49 76 12
29 00.33
57 20.10
29. 00.5322.03 42
37.54 17.44-00 20.00
50. 10.
22.8720.11 41
54. 54.1620 50
63.638.0.37
t0..00
35. 35.9220017
79.41 24024
67.35 +100
92.71#0.37
07.78#1002
NIW
NIW
18 18.9942
36.63
22.
f
1
10
1
10
1
10
10
1
10
10
10
1
10
FL settings
FL settings
f
0.1
1.0
0.1
1.0
f
0.1
1.0
0.1
1.0
100
10
100
10
34

---

Under review as a conference paper at ICLR 2023
FedPop
97. 424
94.98
FedPop
97.91
95.80
FedEM
97 .3.38
95.33
FedEM
97.89
96.46
FedBE
97.16 95.11
FedBE
97. 72.62
96. 33.37
Table 10: (MNIST) Global prediction and personalisation accuracy.
FedPA
93 33.38 94. 94.87
FcdPA
96 66
95.93
(a) Global prediction performance (initial accuracy)
pFedBayes
pFedBayes
97. 97.32
97.88 .7.88
95.49
94.
(b) Personalisation performance
Fed-Prox
97 33
95.28
Fed-Prox
97.90 17.96
96.45
Fed-Avg
97.15 94.86
Fed-Avg
97.83 83
96.10
Fed-BABU
Fed-BABU
97.16 94.59
97.73 95.88
Mix. (K = 2)
Methods
97.94 95.68
Our Methods
98.85
96.67
Mix.
NIW
95.51
NIW
98.78 96.53
97.
settings
1
5
FL settings
5
0.1
0.1
5
8
5
35

---

Under review as a conference paper at ICLR 2023
FedPop
82. 10.40
72.93
FedPop
91.14
86. 86.40
FedEM
83.51
FedEM
92. 02.10
89.65
76.
FedBE
80.35 73.28
FedBE
89. .6.96
85.00 50
(Fashion- MNIST) Global prediction and personalisation accuracy.
FedPA
81.10
73.47
FedPA
91 10.60
88.76
(a) Global prediction performance (initial accuracy)
pFedBayes
pFedBayes
81.20 1..20
72.35
90 01 86.38
(b) Personalisation performance
Fed-Prox
Fed-Prox
83. 73.67
90.72 87. 87.35
Fed-Avg
81.98
73.70
Fed-Avg
90.58 59 86 85.85
Fed-BABU
Fed-BABU
83. 86.86
76.10
91.03 1.3
89.02
Mix. (K = 2)
Mix. =
Table 11:
Methods
84.28
77.60
Our Methods
92.54
89.53
NIW
84.18
77.48
NIW
92.48
89.91
settings
1
5
FL settings
F
1
5
0.1
0.1
5
8
5
36

---

Under review as a conference paper at ICLR 2023
FedPop
85.27
FedPop
88 40
FedEM
85.21
FedEM
88.32
FedBE
85.24
FedBE
88.37
Table 12: (EMNIST) Global prediction and personalisation accuracy.
FedPA
83.20
FedPA
85.10
(a) Global prediction performance (initial accuracy)
pFedBayes
pFedBayes
84.65
$8.12
(b) Personalisation performance
Fed-Prox
Fod-Avg Fed-Prox
85.27
88. 33
Fed-BABU Fed-Avg
85.27
87.92
Fed-BABU
84.33
83.09
Mix. (K 2)
Mix. (K = 2)
Our Our Methods
85.58
Our Methods
88.97
NTW
85.40
NIW
88.84
FL setting
7
1
FL setting
7
f
0.2
0.2
37

---

Under review as a conference paper at ICLR 2023
actually treat 3 (and Φ of feature extractor) random variables, but deterministic variables which are optimized in empirical Bayes learning. This hinders the model from benefiting from hierarchical Bayesian modeling (e.g., they do not have prior distribution p(3) at all); 3) Their optimization is alternating between the feature extractor & and the head prior parameters 3, utterly different from block coordinate optimization alternating between higher level random variables and individual local variables; 4) They did not use variational inference for inference p(=1 Di, Φ, B), but MCMC
sampling (Lagevin dynamics), which is the very reason why they had reduce the size of the latents z' only limited to classification heads, instead of full network parameters as we did.
38


Reviews:
[{"review_id": "hJ8x_T0xj9", "reviewer": "Reviewer_7vmX", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "This paper proposes a novel hierachical Bayesian approach to FL, with variational inference. The idea is to use two levels of random variables, the higher level $\\phi$ which is shared among clients, and lower level $\\theta_i$'s for each client $i$. This allows flexibility of conditional independence and personalization. There are two main methods: Normal-inverse-Wishart (NIW) and Mixture Model (Mix.). The methods are supported by convergence and generalization proofs, as well as experiments on two datasets. \n\n--post rebuttal--\nThanks to the authors for the careful response to my questions. I have updated my score accordingly. After reading the response, I think the paper indeed improves, but there are still some weaknesses to address. It seems that on the theoretical side, the only main novelty is the generalization bound, as the authors admit. However, after checking the proof, I found this generalization bound shares a lot of similarities with Bai et al. 2020. Moreover, in the main paper, Theorem 4.2 needs a lot of discussions, such as the meaning of each term (e.g. $r_n$, $\\epsilon_n$), the implication of this theorem, and comparison with existing results (e.g. Bai et al). On the algorithmic side, since hierarchical Bayesian is proposed before, what is the main novelty of this paper? Is it simply an application for personalized FL? Finally, on the experimental side, after adding more baselines, the improvement over baselines is not significant. For instance, for MNIST the improvement is less than 1%. Even for CIFAR-100, for $\\tau = 10$ the improvement is much less than $\\tau = 1$. This casts doubt on the practicability of the proposed algorithms. ", "summary_of_the_review": "In summary, this paper proposes hierachical Bayesian FL using variational inference which looks interesting. However, the paper largely suffers from presentation, and it is not ready for acceptance till it's more clearly presented.", "strength_and_weaknesses": "Strengths:\n1) The Bayesian approach to FL is interesting and seems suitable for personalization. Using the hierachical approach and variational inference, the authors propose two different choices of hierachical models, that seem to be work well in practice.\n2) The proposed Bayesian FL algorithms are supported by convergence analysis and generalization error bound.\n3) The method is tested on CIFAR-100 and CIFAR-C-100 and compared with baseline algorithms, with noticeable improvement.\n\nWeaknesses:\nAlthough this hierachical Bayesian approach seems interesting, and it is supported by some theoretical analysis as well as experimental comparison, I feel this paper is currently not ready for publication due to several reasons:\n\n1. The motivation is not very clear and some claims are not well-supported. For example:\n 1) Sec 1, paragraph 1: \"FedAvg and FedProx are well-known to suffer convergence issues\". Any reference? What are the assumptions? Why would personalization resolve the convergence issues? \n2) Sec 1, paragraph 2: why would Bayesian perspective be good for personalization? Is the current method complete or principled? Despite this being a bit subjective, I think at least the choices of NIW and Mix. are not principled and ad hoc.\n3) What is the meaning of the higher level random variable $\\phi$? Is it like some prior or shared knowledge? More intuitive explanation is needed.\n4) Figure 1: what are $y^p$ and $x^p$? \n5) Sec 2.1: reference of ELBO and the original ref of block coordinate optimization are needed.\n6) Sec 2.2: the notation $p^p$ seems a bit weird as the two $p$'s have different meanings.\n7) eq.10: what does the approximation sign mean? In which cases does the equality hold?\n8) Sec 3: how the two models are proposed does not seem straightforward. Why are NIW and Mix. good choices and how are they proposed?\n2. Convergence analysis. In Appendix D the authors added 3 assumptions. In Assumption 1, the definitions of \"locally convex\" are not quite clear. Assumptions 2 and 3 have some overlap regarding the Lipschitzness and the radius of the domain needs to be justified.\n1) The convergence of the proposed methods has $O(1/\\sqrt{T})$ convergence rate. This is at the same level as FedAvg (see e.g. [1]). Therefore the proposed method has no advantage over FedAvg which the authors claim to have some convergence issues.\n3. The generalization error bound is not well explained. For example, what is $\\epsilon_n$? What does $\\lambda_i^*$ mean? Why do you use Hellinger distance? How is the error bound compared to (Mohri et al 2019)?\n4. The empirical evalution can be improved. The dataset is mainly CIFAR-100 and its variant, and only 4 baselines are compared. FedAvg and FedProx are pretty similar and more Bayesian approaches should be compared, such as FedPA, FedBE, FedEM, and PredictiveBayes [2].\n\n[1] Adaptive federated optimization. Reddi et al, ICLR 2021. \n[2] Robust One Round Federated Learning with Predictive Space Bayesian Inference, Hasan et al, https://arxiv.org/pdf/2206.09526v1.pdf.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "T7myixvIyMg", "reviewer": "Reviewer_FDT8", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "This paper uses Bayesian paradigm for personalized Federated Learning. \nSpecifically, this paper suggests a two-level hierarchical model for data $D_{1:N}$ distributed on $N$ clients of the form:\n$$\np(D_{1:N} | \\theta_{1:N}) = \\prod_{i=1}^N p(D_i | \\theta_i) \\quad , \\quad p(\\theta_{1:N} | \\phi ) = \\prod_{i=1}^N p(\\theta_i[\\phi) .\n$$\nHere, $\\theta_{1:N}$ represent individual client parameters which are linked trough the globally shared variable  $\\phi$.\nTaking some prior for $\\phi$, the posterior distribution is\n$$\np(\\phi, \\theta_{1:N}| D_{1:N} ) \\propto p(\\phi) \\prod_{i=1}^N p (\\theta_{i} | \\phi) p(D_i | \\theta_i).\n$$\nThis distribution is intractable and the authors use variational inference to approximate it by\n$$\nq(\\phi,\\theta_{1:N} | L_{0:N}) = q(\\phi | L_0) \\prod_{i=1}^N q(\\theta_i|L_i),\n$$\nwhere $L_{0:N}$ is the parameters which are learnt maximizing the ELBO\n$$\n\\mathcal{L} = \\sum_{i=1}^N \\mathcal{L}_i  + \\mathcal{L}_0 ,\n$$\nwhere $\\mathcal{L}_i$ depends on $L_i$ and $D_i$. \n\nTherefore, $\\mathcal{L}$ can be optimized using a block-coordinate optimization scheme. At each step, and given $L_0$, each client optimizes $\\mathcal{L}_i$ and sends its update to a central server which in turn optimizes $\\mathcal{L}_0$ to update the global parameter $\\phi$ for fixed $L\\_{1:N}$. \n\nBased on their inference, the authors also suggest how to make personalization for an unseen new client. \nIn addition, they derive convergence bounds for their FL optimization algorithm and obtain generalization bounds for their variational approximations as the number of observations goes to $\\infty$. Finally, numerical experiments compare the method proposed in the paper with some FL baselines. \n\n", "summary_of_the_review": "The paper is interesting but should be more specific on how it addresses FL problems. In addition, comparisons with other personalization methods should be added. ", "strength_and_weaknesses": "Strength:\n- The paper tackles the important problem of personalization in federated learning using a hierarchical Bayesian approach.\n- The approach using variational inference to infer local and global parameters is interesting.\n- The paper provides theoretical guarantees on the convergence of the method.\n\nWeakness:\n- Practical implementation of the method could be better presented as well as its computational complexity. \nMoreover, the paper would benefit from a better description of how the proposed solution addresses the specific problems of FL, in particular communication constrains and privacy. In the current version, it seems to me that it only focuses on personalization. \n- The method does not address the problem of partial communication which is common in FL.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "RVDn5ILnBa9", "reviewer": "Reviewer_UL1w", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "This paper proposes a hierarchical Bayesian approach to Federated Learning. The proposed Bayesian model makes the block-coordinate descent solution becomes a distributed algorithm. In addition, the paper also derives convergence analysis that shows the block-coordinate FL algorithm converges to a local optimum of the objective at the rate of $O(1/\\sqrt{t})$, which is the same as the regular SGD; and a generalization error analysis that shows the test error is asymptotically optimal. The effectiveness of this approach is demonstrated empirically on two benchmark datasets (CIFAR-100 and CIFAR-C-100). The proposed approach outperforms the best work in terms of both global prediction and personalization. ", "summary_of_the_review": "I would like to recommend the paper for acceptance based on the following reasons:\n1) The paper is clearly written and well organized.  \n2) The proposed method by the paper is well motivated and technically correct. \n3) The effectiveness of the paper is demonstrated both empirically and theoretically. ", "strength_and_weaknesses": "# Strength\n1. The paper is well structured. The paper first introduced a general Bayesian FL framework and derived the ELBO objective optimization function with block-coordinate optimization solution. Then, the paper presents two concreate models (NIW and Gaussian Mixture), which illustrate the hierarchical Bayesian framework very well.\n2. The paper presents principled hierarchical Bayesian approach, in which client’s individual parameters are governed by a common prior. Compare with other Bayesian FL approaches, the proposed method is fully Bayesian and do not need ad-hoc heuristics nor strong assumptions.\n3. The effectiveness of the proposed Bayesian hierarchical FL is empirically tested on benchmark datasets (outperforms the best by large margins most of the time) in both global and personalization prediction.\n4. The paper provides strong theoretical analysis of the convergence and the generalization error of the proposed block-coordinate FL algorithm.    \n\n# Weakness \n1. The discussion with related work needs to include other Bayesian treatment for personalized FL, e.g. [1]\n2. The empirical evaluation is only on two benchmark datasets, considering including more relevant datasets in [2] \n\nReferences: \n[1] Idan Achituve, Aviv Shamsian, Aviv Navon, Gal Chechik, and Ethan Fetaya. Personalized\nfederated learning with gaussian processes. In Advances in Neural Information Processing\nSystems 34: Annual Conference on Neural Information Processing Systems, pages 8392–8406,\n2021.\n[2] Chen, D., Gao, D., Kuang, W., Li, Y., & Ding, B. (2022). pFL-Bench: A Comprehensive Benchmark for Personalized Federated Learning. arXiv preprint arXiv:2206.03655.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "15xmjdSWY8", "reviewer": "Reviewer_SHVY", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "This paper proposes a novel hierarchical Bayesian approach to federated learning, and derive ELBO objective function using variational inference techniques.  Then the block-coordinate descent algorithm is devised, which fit well in the federated learning regime. Theoretical result on the convergence of the algorithm is obtained. Empirical studies demonstrated the superior performance of the proposed method over existing ones.", "summary_of_the_review": "This paper develops a novel federated learning method with hierarchical Bayesian approach, and it can cope with the heterogeneous local distributions. Sound theoretical results on the convergence and the generalization performance are attained. Empirical studies are relatively weak, because only simple dataset, cifar100, is used and the compared algorithms are some standard FL approaches.", "strength_and_weaknesses": "Strength:\n1.\tThe idea of employing Bayesian approach and variational inference techniques for federated learning is of novelty. \n2.\tTheoretical analysis shows the convergence of the proposed algorithm is guaranteed. \n3.\tThe develop framework can tackle the problem of distribution disparity among the local learner. \nWeaknesses:\nThere are some studies that deal with the heterogeneous distribution in domain generalization and adaption, which is not covered by the analysis of related works; the empirical studies also only compare the proposed algorithm with standard FL methods. The conclusions will be more convincing if more comparison is made with some approaches that deal with the heterogeneous distribution.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject