A Bayesian Approach for Personalized Federated Learning in Heterogeneous Settings
Authors: Disha Makhija
Joydeep Ghosh
Nhat Ho
Venue: NeurIPS 2024 poster (2024)
Keywords: Federated Learning
Bayesian Learning
Primary area: deep_learning_architectures
TL;DR: This paper presents a Bayesian approach for personalized FL in heterogeneous settings to enable personalized, privacy-preserving models on decentralized devices, handling challenges like small datasets, model heterogeneity, and privacy constraints.

Abstract: Federated learning (FL), through its privacy-preserving collaborative learning approach, has significantly empowered decentralized devices. However,  constraints in either data and/or computational resources among participating clients introduce several challenges in learning, including the inability to train large model architectures, heightened risks of overfitting, and more. In this work, we present a novel FL framework grounded in Bayesian learning to address these challenges. Our approach involves training personalized Bayesian models at each client tailored to the unique complexities of the clients' datasets and efficiently collaborating across these clients. By leveraging Bayesian neural networks and their uncertainty quantification capabilities, our local training procedure robustly learns from small datasets. And the novel collaboration procedure utilizing priors in the functional (output) space of the networks facilitates collaboration across models of varying sizes, enabling the framework to adapt well in heterogeneous data and computational settings. Furthermore, we present a differentially private version of the algorithm, accompanied by formal differential privacy guarantees that apply without any assumptions on the learning algorithm. Through experiments on popular FL datasets, we demonstrate that our approach outperforms strong baselines in both homogeneous and heterogeneous settings, and under strict privacy constraints.

A Bayesian Approach for Personalized Federated Learning in Heterogeneous Settings
Disha Makhija
Joydeep Ghosh
Electrical and Computer Engineering University of Texas at Austin
Electrical and Computer Engineering University of Texas Austin
Austin, TX 78705
Austin, TX 78705 jghosh@utexas.c edu
dishamutexas.. edu
Nhat Ho
Statisties and Data Science University of Texas at Austin Austin, TX 78705 minhnhat@utexxs. edu
Abstract
Federated learning (FL., through its privacy-preserving collaborative learning ap- proach, has significantly empowered decentralized devices. However, constraints in either data and/or computational resources among participating clients introduce several challenges in learning, including the inability train large model architec- tures, heightened risks of overfitting, and more. this work, we present a novel FL framework grounded in Bayesian learning to address these challenges. Our approach involves training personalized Bayesian models each client tailored the unique complexities of the clients' datasets and efficiently collaborating across these clients. By leveraging Bayesian neural networks and their uncertainty quantification capabilities, our local training procedure robustly Iearns from small datasets. And the novel collaboration procedure utilizing priors the functional (output) space of the networks facilitates collaboration across models of varying sizes, enabling the framework to adapt well in heterogeneous data and compu- tational settings. Furthermore, we present a differentially private version of the algorithm, accompanied by formal differential privacy guarantees that apply with- out any assumptions on the learning algorithm. Through experiments FL. datasets, we demonstrate that our approach outperforms strong baselines both homogeneous and heterogeneous settings, and under strict privacy constraints.
popular
Introduction
Federated Learning (FL) has emerged as pivotal paradigm in various real-world applications, offering decentralized approach that allows participating clients to contribute to a shared model without compromising the privacy of their data. However, implementing FL in practical scenarios poses challenges due to the significant variability among participating clients in terms of their local data and computational resources. Clients with restricted compute capacity may encounter difficulty in training large machine learning models identical to those of other clients, and those with minimal data may struggle obtain reliable estimates of local model parameters.
Due to their ability to generalize under limited data and provide uncertainty quantifications 69 consider Bayesian based learning methods to construct improved local models. Employing Bayesian learning in FL, however, would involve the following steps each client doing local
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

posterior inference to obtain a distribution over weight parameters and then communicating the local posteriors the server, the server receiving the local posteriors from the clients and aggregating them to obtain a global posterior distribution, which is then broadcast to the clients for the next round of training. This entire learning procedure, as turns out, highly resource and communication- intensive. For solving an m-dimensional federated least squares estimation, this method will require (mm)) computation on all the clients and server sites [4] which is much more as opposed to the cost of standard FL (generally (m)).. How could we then utilize the strengths of Bayesian methods for FL settings without paying such high costs? Additionally, given substantial variations in locally available computing resources, how do we still enable efficient learning and collaboration clients? We address these questions by proposing a framework that allows all clients to train their own personal Bayesian models (with varying model complexities), and achieves collaboration across clients by distilling knowledge from the peer clients via a shared unlabelled public dataset and instilling that knowledge in the local models in the form of priors. The challenge of transferring knowledge across models of different architectures is addressed by using the functional (output) space to indirectly determine priors weights of the local model parameters. Furthermore, to prevent the data leaks in FL setup [22] 68 23 and formally guarantee the privacy of the local client data, we present differentially-private version of the algorithm applying formal well-known standard of differential privacy [19], along with privacy analysis and bound on the privacy loss of the entire procedure.
all
This work provides a novel integrated Federated Learning (FL) framework, FedBNN, designed to tackle challenges arising from both limited data and heterogeneous computational resources across clients. Additionally, our method offers valuable characterizations of model uncertainties, and is able to operate under strict data privacy constraints, thereby extending the applicability of FL. crucial domains such as healthcare and legal, where these considerations are paramount. To the best of our knowledge, no previous work has jointly addressed all these learning challenges in the FL context. Our promising results significantly broaden the potential of FL for critical real-world applications. Specifically, our key contributions can be summarized as
We propose a new approach to personalized federated learning utilizing Bayesian princi- ples for improved robustness and reliability, particularly in contexts where data is scarce. Despite its Bayesian framework, this method is designed to be both computationally and communication efficient.
A novel collaboration mechanism based on assigning prior distributions over the model parameters via the output space, instead of directly sharing model parameters' distributions, which can be computationally expensive and raise privacy concerns, is proposed to enable clients having different computational resources to train models with varying complexity. This is important because in real-world FL applications, clients often have vastly different capabilities.
We provide formal differential privacy guarantee for our method that applies to general settings irrespective of the client's 's learning algorithm and show that the method is able to learn effectively even under strict privacy guarantees.
We evaluated our method on several datasets and show that it outperforms the baselines by a significant margin, particularly in heterogeneous data and model settings. This makes FedBNN particularly well-suited for real-world FL. applications, which often exhibit high degrees of heterogeneity.
2 Related Work
This section provides brief overview of the most relevant prior work in the fields of federated learning, Bayesian FL, and Differential Privacy in FL.
Federated Learning FL was introduced as the FedAvg algorithm in the seminal work in 45). Since
then many different modifications have been proposed that tackle specific challenges, including global
FL solutions as well as personalized solutions. FedPD 774, FedSplit (54), and FedDyn [ proposed
methods for finding better fixed-point solutions to the FL optimization problem. [4 73 66
58
[S
show that point-wise aggregate of the local client models does not produce a good global model
and propose alternate aggregation mechanisms to achieve collaboration. Personalized FL has been
2

---

Table 1: Contrasting our method, FedBNN, against previous works.
Addressed Challenges
Method
Limited
Heterogeneous
Uncertainty Quantification
Privacy
Compute
FedProx
pFedME
FOLA
PFedGP
PFedBayes
FodPop
/
FedAUX
FedBNN
approached in many ways like meta-learning 20 8 31 33 multi-task learning (59 38 60], by clustering the clients [56] 26 and others [7
uses Bayesian view based analysis to obtain a better trade-off between personal and global models. Knowledge distillation for personalized FL settings has also been used previously for training heterogeneous models in 40 (0)] Several other methods have proposed enhancements in FL learning and privacy by using an auxiliary dataset like [3 55 36 52). Most of these methods rely on the well-established knowledge distillation procedures. But since the transfer of knowledge or information between Bayesian models itself has remained inadequately addressed, these methods are not easily extensible to Bayesian settings. Our method the other hand, utilizes a novel method that enables collaboration across client specific Bayesian models by transferring knowledge through prior specification mechanism in the output space, which also enhances the field of Bayesian knowledge distillation.
non-Bayesian settings 336
Bayesian Federated Learning Bayesian approaches for federated learning can also be broadly divided as methods using Bayesian inference to obtain a global model and personalized Bayesian learning methods. Amongst the methods that train global model, some methods just use Bayesian mechanisms for achieving collaboration among non-Bayesian local models, like FedBE [I5] which uses Bayesian mechanism to aggregate the locally trained neural networks to obtain Bayesian ensemble at the server, [9] which suggests using an MCMC based method for obtaining a global model from the local models, and PFNM (73) and FedMA (66) which use
process to obtain the global models. Other methods that train local Bayesian models at the clients and at the server include FedPA [4] that uses Laplace approximations for an efficient way of computing local and global posteriors,
18] that suggests the use of Bayesian Optimization and Thompson Sampling to obtain the solution to the global optimization problem, recently, 49 did empirical study on various ways of aggregation mechanisms for local variational Bayesian neural networks and their effects on the solution. These methods that focus on obtaining global solution less suited for the statistical heterogeneity present across clients 14, and therefore we focus more on the methods that build personalized Bayesian solutions for clients. Among such methods, pFedGP (3 is a Gaussian Process based estimation method that utilizes Deep Kernel Learning to collaboratively train a single deep neural network with FedAvg and then uses personalized GPs for prediction. FedLoe [70] also uses GP in FL but for regression tasks, PFedBayes (75] uses variational inference locally to optimize a loss at each client that is combination of the data likelihood term and distance the prior and iteratively determines the prior from the global posterior distribution. FOLA [1 proposed using Laplace Approximation for posterior inference at both the server side and the client side, PAC-FL [I]] and 34 65 51] also proposed variants of methods that assume Bayesian models on local clients but for all of them main assumption is that the local model parameters are generated from a shared global distribution thus making them useful only in homogeneous settings. All the methods described above choose priors by assuming a distribution over values for each weight, and thus choosing an appropriate and meaningful prior becomes a challenge [14). These issues led us to use functional space priors instead which have been explored in limited centralized settings [63] 6 21 but not in FL. But most importantly, none of these methods are designed or could be easily extended to work with compute heterogeneous settings limiting the applicability of these solutions in several real-world scenarios. Table [ compares our approach with the most closely related works.
Differential Privacy in FL Since decentralized learning does not guarantee that the data will remain private, it is important that formal rigorous guarantee be given on the data that is leaked by the algorithm. Seminal works in DP propose using a Gaussian noise mechanism by adding Gaussian noise to the intermediate results and achieving a bound on the algorithm by using composition
3

---

results [9]46]
44 independently proposed DP-FedSGD and DP- algorithms, which enhance FedAvg by adding Gaussian noise to the local client updates. Several other works focus analyzing the privacy-utility trade-off in DP in FL setting 225 27
46 32]. For FL. 124 and
S 64 39
Recently, 30 proposed a DP-based solution for personalized FI. that works only for linear models.
And then 47 improved it for general models and heterogeneous data in FL.. These methods, however,
mostly focus on privacy guarantees while solving the non- con-Bayesian FL optimization problem.
3
Methodology
In this section, first go the problem setting and background, and then present proposed framework, FedBNN, with details of all the key components.
3.1 Background
Problem Description Consider an FL setting with N clients where each client i has local dataset X' of size n, drawn from the local data distribution D;- The goal of personalized federated learning procedure is to obtain optimal weights for each client's local model, WT, given the entire data, X = - X3 through collaboration but without compromising client data privacy. However, the learning procedure faces challenges that are posed due system heterogeneity and statistical heterogeneity. System heterogeneity refers to the variable amount of data and compute resources across clients, meaning. i) the data resources on each client vary widely, i.c., R& >> ny for some clients k and I, and ii) the compute across clients is non-identical due to which it is not possible to train models uniform architectures across clients, leading to non-identical weights, i.e., W; + W; for different clients i and j. Statistical heterogeneity implies that the data distribution across clients is non-IID.
Bayesian Learning Instead of obtaining the optimal values of the model parameters, W,, Bayesian learning aims to learn posterior distributions (probability distributions over the values) for all the model parameters from the given data PPW](X) Thus, in personalized Bayesian FL procedure, the modified goal would be to learn distributions for local weights, (P) from = - x-- However, the exact inference for obtaining the posterior distribution for each of the weight parameter in the network is intractable and several approximations have been studied to obtain approximate distributions. Variational inference [29] an approximation method that tries to learn parameterized distribution q (0) from a family of distributions Q, typically of simpler form, by optimizing the parameters such that the new distribution g(w) (0*). obtained for the optimal value of 0, close to the desired posterior distribution IP()))). Precisely, 0" is obtained by solving the following optimization problem and its expansion given below
8" = arg min [[[[W))][P(())))) 0
arg min KL(g(W10)\p<(W:e)) E,(w1o) (o 0
and then q(w (*)) is used in place of IP) The optimization objective minimizes the distance of g(Wl)) prior distribution p(W: used encode any prior information about the parameters, while also maximizing the likelihood of the observed data X' under q(W)l)) A more detailed discussion of Bayesian learning is included in Appendix A Even though Bayesian approaches are more computationally expensive than their point-estimation counterparts, their superior capabilities for uncertainty quantification and performance in small data settings outweigh the extra compute costs in many critical applications. Moreover, recent innovations like Bayes by Backprop [10] which carefully uses the backpropagated gradients for learning the parameters of posterior distributions, drastically reducing the added computation costs.
3.2 FedBNN Methodology
The FedBNN framework works iteratively in two steps local optimization on the individual clients to obtain local posterior distributions over the model parameters, and a global collaboration step

---

where the output from each client is appropriately aggregated at the server and broadcast to all the clients for the next rounds of training. These two steps are further described below, and the detailed algorithm and the overview diagram are included in the Appendix B in Algorithm land Figure respectively.
Local Setting Let each client in the network be training a personalized Bayesian NN, which for the client i is denoted by 4; and is parameterised by weights W. As commonly used in the literature, we assume that the individual weights of the BNN are Normally distributed and satisfy mean-field decomposition, i.c., Wa ~ (N o.aa) for a E [_,..., [W.I] where Hi,a is the mean of the Gaussian distribution for the parameter Gaussian distribution for the same parameter. To guarantee that dta. takes non-negative values for all clients and all parameters a, we use technique commonly used in inference procedures (TO), and replace each Ti,a by another parameter Pi.o during the training, with dia log(1 + exp(pt.a)). The individual weights of the local BNN, Wi, are also assumed to each have Gaussian prior distribution, p( a( )),a), parameterized by vii,a = (1% P.a)-
the 5th client and of is the variance of the
3.2.1 Global Collaboration
We attain collaboration amongst clients via auxiliary dataset called the Alignment Dataset (AD). This is an unlabeled dataset typically small in size, and is used for providing peer supervision to the individual clients by helping clients distill knowledge from other peer clients without explicitly sharing a large number of locally learned parameter weight distributions. The experiments in Figure and Table show the effect of the varying size and distribution of AD in achieving effective collaboration.
In heterogeneous settings, the use of non-identical architecture models (W; * W3) means that there no direct way of aggregating the distributions for prior specification. In fact, even in homogeneous settings, aggregating the weight distributions can be prone to errors due to reasons like insufficient understanding of the weight space, non-alignment of weights across models, etc. Thus, for the purpose of collaboration, we use the function-space of the networks rather than the weight space. Specifically, in each global communication round, the server shares the AD with all the clients. The clients do a forward pass on AD to obtain the local output $, (AD), where the local output of the
7th client is approximated by drawing m sets of weight samples, wi(") : E [1, m). from its local posterior distribution (() using Monte Carlo sampling and aggregating the outputs under each
of these samples &;(AD) 1 Σ (AD; w!( The obtained output for AD on each client is m
then sent back to server which forms an aggregated representation, denoted by (AD). obtained via a weighted aggregation of all clients' outputs, i.e., (xx) = Σ (X). By default, all weights are considered the same, however the formulation provides flexibility, for example to accommodate situations where the aggregation weights could represent the relative strength of each client in terms of its data or compute resources, i.c., clients with high compute (or data) resources receive more weight as compared clients with lower amount of resources. The obtained & (AD) is then uploaded to all the clients for use in the next round of local training. More details about the Aligmment Dataset (AD) along with the explanations and experiments on the size, distribution, availability etc. of the AD are included in the Appendix E
3.2.2 Local Optimization on Clients
Prior Specification Design The Bayesian framework provides a natural way of incorporating supervision in the form of priors. Conventional methods in Bayesian deep learning provide direct priors for model weights as distribution over values. However, the relationship between the values of the model weights/paramettes and the outputs is complex and the priors in model's weight-space do not directly capture the desired functional properties. Also, since the number of parameters in a neural network large, most prior specifications tend to take a simplistic form like an isotropic Gaussian, to make inference feasible. Thus, learning by specifying prior distributions over weights does not always help translate prior knowledge the learning process. In this work, we consider a way of specifying priors in the functional space by first optimising the Bayesian neural networks over the prior parameters for a fixed number of steps so that the BNN achieves a desired functional output. These intuitive priors help in explicitly instilling the external knowledge during the training of the
5

---

neural networks. Let ;W;;; v represent the prior function over the weights W; and is parameterized by v. with & = [(uf,0 L.)) a E ....... [Will, the prior parameters that determine the prior distributions are learned by solving an optimization problem as below:
arg min m 2,D,:; W.)),
where d is a suitable distance function and Y represents the desired output, resulting in optimal priors p(W; w²). Below we provide details of the prior specification for our method.
Local Optimization For the local optimization, the individual clients learn P(WA)&;; via vari- ational inference. As described above, variational learning algorithm tries to find optimal pa- rameters 0" of parameterized distribution g(W)]00 among family of distributions denoted by Q. our setting, we set the family of distributions, 2, to be containing distributions of the form Wi, N (( of.a) for each parameter Wi, for a E []...., W&I]. For inference in Bayesian neu- ral networks, we use Bayes by Backprop [10] method to solve the variational inference optimization problem.
At the beginning of each local optimization procedure (in each global communication round specific client is selected), we use the global information obtained from the server ((AD) to intialize the prior for the BNN. Specifically, the beginning of each local training round, the selected clients first tune their priors to minimize the distance between the local output, (AD:; W;) and the aggregated output obtained from the server, (AD). Since the aggregated output represents the collective knowledge of all the clients and may not be stricily precise for the local model optimization, we consider this aggregated output as "noisy" and correct it before using for optimization. Specifically, we generate @poorecied as convex combination of the global output and the local output for a tunable parameter y. For the ith client,
S conrected = >DDA + (L, W.).
The prior optimization steps then optimize the distance between dromecied and D; (AD; W,) to train the prior parameters vh, with the aim of transferring the global knowledge encoded in dpcornetedd to the local model. Precisely,
arg min di ((.. @purrctedd (,(DD; Wc).
When the outputs 4 (X; W) are logits, we use cross-entropy or the negative log lilkiiiikooodd loss as the distance measure. The optimization involves training the client's personal BNN o to only learn the parameters of the prior distribution denoted by tb. This way of initializing the BNN prior enables translating the functional properties, as captured by <ADD W;). to weight-space distributions. The optimal prior parameters are then kept fixed while training the BNN over the local dataset. The local optimization procedure now works find the best ((W, fixing the prior distribution through the following optimization problem:
arg min (L(g((A))))PP(M) +f) 100 (og 10)
(5)
which is similar to the optimization problem defined in Equation exxcept that now the prior parameters are optimized so that the obtained prior distributions capture the global knowledge and can guide the local learning process to make q(W,10) close to the global collective knowledge.
3.2.3 Achieving Differential Privacy
In this variation, to control the release of information from the clients, we add a carefully designed Gaussian mechanism wherein we add Gaussian noise to the & (AD) that being shared by each client. Specificallyy each client uploads (AD) = (AD) +NO, oz) to the server and then the server aggregates (AD)) across clients to obtain and broadcast ((AD) which is used by the clients in their next round of local optimization. The variance of the noise depends on the required privacy guarantee.

---

4 Privacy Analysis
Though our algorithm is inherently quite private as it refrains from explicitly sharing model weights, we can also provide a formal Differential Privacy based guarantee. Our analysis in this section focuses on providing record-level DP guarantee over the entire dataset Y. This analysis quantifies the level of privacy achieved towards any third party and hotest-but-curious server. In this section we directly present the key result of our analysis. Due to the lack of space, additional definitions, results and the proof for the theorem mentioned in Appendix C
Theorem 4.1 (Privacy Budget). The proposed algorithm is (c, 8) -)-ifferrnttilly private, if the total privacy budget per global communication round per query is set to
e
4FKlog
for E number of global communication rounds and K number of queries to the algorithm per round.
The parameter is related to the Gaussian noise by p 20 The detailed proof is included in Appendix C Our analysis does not assume any specifics of how each client is trained and is therefore applicable more general settings. Note that we present pessimistic analysis by providing worst- case analytical bound, wherein we assume that a change in single data point may entirely change the output of the algorithm, and also since the public dataset remains common throughout the rounds, the actual privacy loss due to querying on the public dataset does not typically add up linearly. Yet the above analysis shows that we have several knobs to control to achieve the desired privacy-utility trade off balancing the number of global communication rounds with local epochs, reducing the number of queries, and the standard noise scale. By appropriately tuning these controls we are able to achieve good performance with a single digit ( 9.98) privacy guarantee and = 10-4
5 Experiments
In this section, we present an experimental evaluation of our method and compare it with different baselines under diverse homogeneous and heterogeneous client settings. Specifically, experiment with three types of heterogeneity -i) heterogeneity in data resources (amount of data), ii) heterogeneity in compute resources, and iii) statistical heterogeneity (non-III data distribution across clients). We
also discuss the change in performance of our method when the degree and type of heterogeneity changes. Due to the space constraint, additional experiments on varying the size and distribution of the AD, privacy-utility trade-off and model calibration are included in the Appendix E G and D respectively.
5.1 Experimental Details
Datasets We choose three different datasets commonly used in prior federated learning works from the popular FL benchmark, LEAF [33 including MNIST, CIFAR- 10 and CIFAR-100. 100. MNIST contains 10 different classes corresponding to the 10 digits with 50,000 28 28 black and white train images and 10,000 images for validation. CIFAR- 10 and CIFAR- contain 50,000 train and 10,000 test-colored images for 10 classes and 100 classes respectively. The choice of these datasets is primarily motivated by their use in the baseline methods.
We simulate three different types of heterogeneous settings - corresponding to heterogeneity compute resources, data resources and the statistical data distribution. Before starting the training process, we create N different clients with different compute resources by randomly selecting fraction of clients that represent clients with smaller compute. Since these clients do not have large memory and compute capacity, we assume that these clients train smaller-size BNNs as opposed the other high-capacitt clients that train larger VGG-based models. In particular, the small BNNs were constructed to have either 2 or 3 convolution layers, each followed by ReLU and fully-connected layers at the end, and a VGG9-based architecture was used for larger BNNs. The number of parameters in smaller networks is around 50K and that in larger networks is around 3M. Since the baselines only operate with identical model architectures across clients, we use the larger VGG9-based models on the basclines for fair comparison. We include the results of our
Simulation Details
7

---

method in both homogeneous compute settings (similar to baselines) as well as in heterogeneous compute settings wherein we assume that 30% of the total clients have smaller compute and are training smaller-sized models.
Next, we also vary the data resources across clients and test the methods under 3 different data settings small, medium and full. The small setting corresponds to each client having only 50 training data instances per class, for the medium and full settings each client has 100 data instances and all available data instances per class respectively for training. We simulate statistical heterogeneity by creating non-IlD data partitions across clients. We work in a rather strict non-IlI setting by assuming clients have access to data of disjoint classes, For each client a fraction of instance classes is sampled and then instances corresponding to the selected classes are divided amongst the specific clients. For the included experiments, we set number of clients N = 20 and divide the instances on clients such that each client has access to only of the 10 classes for MNIST and CIFAR- 10, and 20 out of 100 classes for 100.
Table 2: Test accuracy comparsion with baselines in non-IlI settings.
MNIST
CIFAR 10
CIFAR 100
Method
(small)
(medium)
(full)
(small)
(medium)
(full)
(small)
(medium)
(full)
(Non-Bayesian)
Local Training
88.7 + 1.2
90.1 + 1.0
91.9 + 1.1
$3.9 + 2.1
59.5 + 1.8
70.8 + 1.4
28.8 + 1.8
32.7 + 1.9
43.5 + 1.6
FedAvg
88.2 0.5
90.15 122
92.23 1:23110 1.0
43.14 12
56.27 1.8
78.17
32811 2.
1.6
36.3 6633:...
FedProx
86.9 -: 0.8
89.91 - 0.7
93.1
44.27
12
1.
0.6
28
34.31
1.4
37.
FedAUX
90.1 1.6
1.34
94.4
60.01
6.76
77.0
37.05
4.
452
pledME
91.95 2.1
93.39 = 1.2
0.5
4.
64.57
12
32
2.
non-Bayesian KD
89.1 - 0.4
92.5 0.2
93.2 + 0.3
33,9 1.3
5.22
69.8
1.0
26.1
35.2 1.2
42.7 1 0.8
(Bayesian with Homogeneous Architectures)
pFedGP
86.15 + 1.3 90.59 = 1.7
94.92 0.3
45.62
5..24
0.7
47.06
1.3
1.
54.54
0.2
PFedBayes
94.0 = 0.2
94.6 :
95.5 + 0.3
58.7
6.
78.3
0.5
39.51
41.43
0.4
47.67
+ 1.1
FOLA
91.74 1.0
92.87 1 0.8
95.12
0.6
43.29
1 0.9
45.94
1.7
67.98
1 0.5
33.42
1
1.3
48.8
1.812
43.2
1.6
Ours (olomo)
94.9 1 1.0
95.72 1 0.8
96.21
0.3
70.6
1 111
72.3
0.6
79.7
1 0.3
49.65
1 1.4
55.4 1 0.8
57.3
0.8
Ours (Ieetero)
93.1 1 1.1
94.4
1
95.9
1
0.2
68.17
2.0
71.73
113
78.7
1 0.7
47.5
49.10 1.1
51.1
1.1100 0.7
Ours (MeeteooDP))
89.82 2.3
90.21 1 1.6
91.43
1 14
$4.9 1
1.91
61.83
1.4
74.3
1
1.6
43.7
1.3
44.5
47.0 1.5
(DP-Baseline) DP-FedAvg
80.1 + 1.7
85.2 1.8
86.2 1.7
35.17 0.8
50.22 1.1
74.6 122 26.5 0.3
30.7 1.4
0.6
Training parameters and Evaluation We run all the algorithms for 200 global communication rounds and report the accuracy on the test dataset at the end of the 200 round. The number of local epochs is set to 20 and the size of AD is kept as 2000, Each client is allowed to train its personal model for fixed number of epochs, which is kept to 50 in experiments. before entering the collaboration phase. The hyyer-parametent of the training procedure are tuned on a set-aside validation set. At the beginning of each global communication round, for optimizing the prior parameters at each client according to Equation 4 we use an Adam optimizer with learning rate=0.0000 0001 and run the prior optimization procedure for 100 steps. Then with the optimized prior we train the local BNN using with Adam optimizer, learning rate 0.001 and batch size 128. The noise effect y is selected after fine-tuning and kept to be 0.7. 7. For these experiments, the aggregation weight Wy for each client.) used to compute (X) is set to 1/N, and the AD is obtained by using a separated subset of the dataset in considcration. All the models are trained on a GPU machine with GeForce RTX 3090 GPUs and 24GB per GPU memory. For evaluation, report the classification accuracy obtained by running the trained models on test datasets from the MNIST, CIFARIO and CIFAR100 datasets.
Bayes-by-Backkpeg
Baselines We compare our method against the standard non- non-Bayesian FL algorithms and Bayesian-Fl methods that build personalized models for clients. We also show results of differentially private FedAvg algorithm under similar privacy guarantee to provide perspective on the privacy. Apart from local training where all clients train independent models locally without collaboration, the non-Bayesian FL baselines include i) FedAvg, ii) FedProx, iii) pFedME (which uses personalized models each client using Monreau envelopes in loss). We also compare our method to other
8

---

baselines that use an auxiliary dataset for collaboration in non -oyysiani FL, namely i) FedAUX, which uses federated distillation for achieving collaboration in FL, we do not use FedMD [36] as a baseline since it requires labelled auxiliary dataset. Further, we create baseline corresponding to the non- non-Bayesian version of our method that works with knowledge distillation and call it non- Bayesian KD. The Bayesian FL baselines include i) pFedGP, Gaussian process based approach that trains common deep kernels across clients and personal tree-based GPs for classification, ii pFedBayes. which uses a variational inference-based approach for personalized FL by training personal models which are close to the aggregated global models, iii FOLA, bayesian method using Gaussian product for model aggregation. And lastly, the DP baseline includes i) Dv. the FedAvg algorithm with gradient clipping and noise addition to the gradient at each client. The size of the AD is changed to 1000, number of local epochs to 40 and global communication rounds to 100 for DP-based experiments. For all the experiments, the paraa--arameters were obtained by tuning on held-out validation dataset. We used own implementation of the pFedBayes algorithm since the source code was not publicly available but we could not compare against FedPop due to the lack of some implementation details and publicly unavailable code.
60-
Pataro-Fadiaaae
pFacllasn
62.5
60.0
pReeCr
peedBayes
75
rabers FedBlayas
pladiayyas
31.5-
Toe: ACOUTIC
60-
00-
45.0
9.
3.2
0.4
0.5
0.5
0.
24
1.0
0.0
0.2
D4
0.9
C.8
3.0
Pradton of clients with sma compute capecity (1c1
Pradton of Ehenta win small compate capucity I8D1
Practie of clintt with small collnet Int
(a)
(b)
(c)
Figure 1: Performance comparison of our method with baselines under different types and varying degree of heterogeneity for CIFAR-10 dataset with 20 clients. Figure (a) is for heterogeneity in compute capacity across clients under non-IID data setting, figure (b) for compute heterogeneity under IID setting, and figure (c) for heterogeneity in data resources. When fraction of clients in the setting have low computing resources, the baselines being homogeneous can only train smaller models all the clients as shown by constant performance. The results show that our method is more tolerant to both model heterogeneity and data heterogeneity across clients.
5.2 Results
The performance of our method and the baselines under the non-III data setting are reported in Table 2 Under the non-IID setting, report the results corresponding to different dataset sizes each client. To recall, in the small, medium, and full settings, each client has access to 50, 100, and all training data points per class respectively. We observe that our method with homogencous architectures across clients outperforms all other baselines. Moreover, when we consider the performance of our method under a heterogeneous setting by considering 30% of the total clients to be small capacity, it is evident that our method is better than the higher capacity homogeneous baselines for more complex tasks like in CIFAR-10 and CIFAR-100. On average, our method achieves about 6% performance improvement over the baselines in the small and medium data settings. Figure I compares the performance of our method with the gighest-perfomiing baselines under model, data and statistical types of heterogeneity. Since our method can work with heterogeneous clients, we see that just by the proposed collaboration and having higher capacity clients in the FL ecosystem, the lower capacity clients are able gain about 10% increase in their performance. Also, the performance degradation of our method with change in the number of clients with limited data resources more graceful as compared to the baselines. In an additional experiment intended to compare the performance of the baseline methods with additional data, we trained the priors for baseline methods' encoders using the unlabeled data, AD, before starting their own prescribed FL procedure. We observed that the performance of the baseline methods does not change on doing this because the FL procedure that they incorporate forgets all the prior existing local knowledge at the client side. A similar result was also reported in (55). The superior performance of our method could be attributed to the innovative and effective collaboration achieved by first distilling peer knowledge in the form of the aggregated output on the

---

AD, and then ensuring that this knowledge is successfully transferred to each client by specifying priors in the sunctional-space of the client model. Furthermore, the parameter in Equation Caaatioon@alloyyy Jallows the clients the flexibility to choose the amount of global knowledge that needs to be incorporated, providing flexibility on the degree of personalization.
Discussion
This paper introduced novel method for personalized Bayesian learning in heterogeneous FL settings and demonstrated that it able to outperform existing approaches under different types of heterogeneous situations, while also providing privacy guarantee and calibrated responses. The experiments show that the method is particularly useful for clients with lower data and lower compute resources they benefit the most by the presence of other, more powerful clients in the ecosystem. While our method assumes the availability of a small, unlabelled auxiliary dataset at the server, is typically a very mild requirement such data often be obtained from several open sources on the web. In many cross-silo and cross-device applications, the server often possesses its dataset alongside private data from elients. For example. hospitals with access patient records may combine this data with private patient data collected from individual devices such as wearables or sensors for FL, source code generation applications might leverage open-source code along with private code repositories from developers, etc. [ also mention use-cases where such data is available in real-world. Recent advances in generative AI have made creation of synthetic data for training much easier task. The privacy analysis on the method provides an intuitive and rigorous guarantee with various tunable knobs that can be adjusted to achieve the desired privacy- cutiliyy trade-off. And while the application explored in the proposed work consists of image related tasks, both the proposed framework and the privacy analysis are generic and independent of specific training algorithms, therefore resulting in its wide applicability in various applications across data modalities. Also, while Bayesian methods are inherently more computationally expensive as they have to maintain distributions rather than point estimates, this extra work is invaluable in many applications where uncertainty quantification is important, for example to help engineers account for uncertaintics in material properties, loading conditions, and manufacturing processes, leading safer and more reliable designs (53). The recent use of transformer based Bayesian methods (76] 42 62 in varied applications indicate that the proposed framework can be also applied to settings where much larger neural networks required. One limitation that originates from the Bayesian nature, and is common to all applications of Bayesian learning, is that the exact inference of the posterior distributions is infeasible and therefore variational approximation has been used for inference of the posterior distributions.
10

---

References
[1] Durmus Alp Emre Acar. Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. In International Conference on Learning Representations, 2021.
[2] Idan Achituve, Aviv Navon, Yochai Yemini, Gal Chechik, and Ethan Fetaya. GreTree A Gaussian process classifier for few-shot incremental learning. In Proceedings of the 38th International Conference on Machine Learning, pages 54-65. PMLR, 2021.
[3] Idan Achituve, Aviv Shamsian, Aviv Navon, Gal Chechik, and Ethan Fetaya. Personalized federated learning with gaussian processes. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[4] Maruan Al-Shedivat, Jennifer Gillenwater, Eric Xing, and Afshin Rostamizadch. Fedcrated learning via posterior averaging: A new perspective and practical algorithms, 2021.
[5] Borja Balle, James Bell, Adrià Gascón, and Kobbi Nissim. The privacy blanket of the shufile model. In Advances in Cryptology CRYPTO 2019, pages 638-667. Springer International Publishing, 2019.
[6] Racf Bassily, Albert Cheu, Shay Moran, Aleksandar Nikolov, Jonathan Ullman, and Steven Wu. Private query release assisted by public data. In Hal Daumé IIII and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 695-703. PMLR, 13-18 Jul 2020.
[7] Raef Bassily, Shay Moran, and Anupama Nandi. Learning from mixtures of private and public populations. In Advances Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[8] Martin Beaussart, Felix Grimberg. Mary-Anne Hartley, and Martin Jaggi. WAFFLE: weighted averaging for personalized federated learning. CoRR, ahs/21 10.06978, 2021.
[9] Shrey Bhatt, Aishwarya Gupta, and Piyush Rai. Bayesian federated learning via predictive distribution distillation, 2022.
[10] Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In Francis Bach and David Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 1613-1622. PMLR, 2015.
[11] Mahrokh Ghoddousi Boroujeni, Andreas Krause, and Giancarlo Ferrari Trecate. Personalized federated learning of probabilistic models: pac-hayesian approach, 2024.
[12] Mark Bun and Thomas Steinke. Concentrated differential privacy: Simplifications, extensions, and lower bounds. In Martin Hirt and Adam Smith, editors, Theory of Cryptography, pages (35-658. Springer Berlin Heidelberg, 2016.
[13] Scbastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečný, H. Brendan McMahan, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for foderated settings, 2019.
[14] Longbing Cao, Hui Chen, Xuhui Fan, Joao Gama, Yew-Soon Ong, and Vipin Kumar. Bayesian federated learning: A survey, 2023.
[15] Hong- You Chen and Wei-Lun Chao. Fed(be): Making bayesian model ensemble applicable to federated leamning In International Conference on Learning Representations, 2021.
[16] Huili Chen, Jie Ding, Eric Tramel, Shuang Wu, Anit Kumar Sahu, Salman Avestimehr, and Tao Zhang. Self-aware personalized federated learning. 2022.
[17] Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated leamnin. In Marina Mcila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 2089-2099. PMLR, 18-24 Jul 2021.
[18] Zhongxiang Dai, Kian Hsiang Low, and Patrick Jaillet. Federated bayesian optimization via thompson sampling. CoRR, 0. 101544, 2020.
[19] Cynthia Dwork and Aaron Roth. The algorithmic foundations of differential privacy. Found. Trends Theor. Comput. Sci., 9(3 (33-4):211-221 407, 2014.

---

[20] Alireza Fallah, Aryan Mokhtari, and Asuman E. Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin, editors, Advances in Neural Information Prucessing Systems 33; Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[21] Daniel Flam-Shepherd. Mapping gaussian process priors to bayesian neural networks. 2017.
[22] Liam Fowl, Jonas Geiping, Wojeiech Czaja, Micah Goldblum, and Tom Goldstein. Rohhing the fed: Directly obtaining private data in federated learning with modified models. ArXiv, abs/21 10.13057, 2021.
[23] Liam H Fowl, Jonas Geiping, Steven Reich, Yuxin Wen, Wojciech Czaja, Micah Goldblum, and Tom Goldstein. Decepticons: Corrupted transformers breach privacy in federated learning for language models. In The Eleventh International Conference on Learning Representations, 2023.
[24] R. C. Geyer, T. Klein, and M. Nabi. Differentially Private Federated Learning: A Client Level Perspective. ArXiv e-prints, 2017.
[25] Badih Ghazi, R. Pagh, and Ameya Velingker. Scalable and differentially private distributed aggregation in the shuffled model. ArXiv, has/1906.083220 2019.
[26] Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran, An efficient framework for clustered federated learning. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-, 2020, virtual, 2020.
[27] Antonious Girgis, Deepesh Data, Suhas Diggavi, Peter Kairouz, and Ananda Theertha Suresh. Shuffled model of diffcrential privacy in federated learning. In Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pages 2521-2529. PMLR, 2021.
[28] José Miguel Hemanddz-LLobato and Ryan P. Adams. Prohabilistic backpropagation for scalable learning of bayesian neural networks. In Proceedings of the 32nd International Conference on International Conference on Machine Learning- Volume 37, ICML' 15, page 1861-1869. JMLR.org, 2015.
[29] Geoffrey E. Hinton and Drew van Camp. Keeping the neural networks simple by minimizing the description length of the weights. In Proceedings of the Sixth Annual Conference on Computational Learning Theory. COLT '93, page 5-13. Association for Computing Machinery, 1993.
[30] Rui Hu, Yuanxiong Guo, Hongning Li, Qinggi Pei, and Yanmin Gong. Personalized federated learning with differential privacy. IEEF Internet of Things Journal, 7(10)9530-9539 2020.
[31] Yihan Jiang, Jakub Konečný, Keith Rush, and Sreeram Kannan. Improving federated lcarning personaliza- tion via model agnostic meta learning. CoRR, abs/1909. 12488, 2019.
[32] Peter Kairouz, Sewoong Oh, and Pramod Viswanath. The composition theorem for differential privacy. In Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 1376-1385. PMLR, 2015.
[33] Mikhail Khodak, Maria-Florina Balcan, and Ameet S. Talwalkar. Adaptive gradient- based meta-learning methods. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer. Florence d' Alchh---Be Buc, Emily B. Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference Neural Information Processing Systems 2019, NeurIPS 2019, December 8- 14, 2019, Vancouver, BC, Canada, pages 5915-5925 2019.
[34] Nikita Yurevich Kotelevskii, Maxime Vono, Alain Durmus, and Eric Moulines. Fedpop: A bayesian approach for personalised federated learning. In Alice H. Oh, Alckh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[35] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan. and R. Garnett, editors, Advances in Neural Information Processing Systems. Curran Associates, Inc., 2017.
[36] Daliang Li and Junpu Wang. Fedmd: Heterogenous federated learning via model distillation, 2019.
[37] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar and Virginia Smith. Federated optimization in heterogeneous networks. In Inderjit S. Dhillon, Dimitris S. Papailiopoulos, and Vivienne Sze. editors, Proceedings of Machine Learning and Systems 2020, MLSys 2020, Austin TX. USA, March 2-4, 2020. misys.org, 2020.
12

---

[38] Tian Li. Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 Proceedings of Machine Learning Research, pages 6357 6368, PMLR, 2021.
[39] Yiwei Li. Tsung- Hui Chang, and Chong- Yung Chi. Secure federated averaging algorithm with differential privacy. In 2020 IEEE 30th International Workshop on Machine Learning for Signal Processing (MLSP). pages 1-6, 2020.
[40] Tao Lin, Lingjing Kong. Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. In H. Larochelle, M. Ranzato, Hadsell, M. F. Balcan, and Lin, editors, Advances in Neural Information Processing Syxtems, volume 33, pages 2351-2363, Curran Associates, Inc., 2020.
[41] Liang Liu, Feng Zheng, Hong Chen, Guo-Jun Qi, Heng Huang, and Ling Shao. A bayesian federated learning framework with online laplace approximation. 2021.
[42] Ahmed Maged and Min Xie. Uncertainty utilization in fault detection using bayesian deep learning. Journal of Manufacturing Systems, 64:316-329, 2022.
[43] Disha Makhija, Xing Han, Nhat Ho, and Joydeep Ghosh. Architecture agnostic federated learning for neural networks. International Conference on Machine Learning, 2022.
[44] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-Efficine Learning of Deep Networks from Decentralized Data. In Aarti Singh and Jerry Zhu, editors, Proceedings of the 20th International Conference on Artificial Intelligence and Statis- tics, volume 54 of Proceedings of Machine Learning Research, pages 1273 1722-222. PMLR, 20 22 Apr 2017.
[45] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication---ffiice Learning of Deep Networks from Decentralized Data. In Aarti Singh and Jerry Zhu, editors, Proceedings of the 20th International Conference on Artificial Intelligence and Statis- tics, volume 54 of Proceedings of Muchine Learning Research, pages 1273-1282. PMLR, 20-22 Apr 2017.
[46] Tlya Mironov. Rényi differential privacy. In IEEF IEEE Computer Security Foundations Symposium (CSF), pages 263-275, 2017.
[47] Maxence Noble, Aurélien Bellet, and Aymeric Dieuleveut. Differentially private federated learning on heterogeneous data. In Proceedings of The 25th International Conference on Artificial Intelligence and Stutistics, volume 151 of Proceedings of Machine Learning Research, pages 10110-10445. PMLR, 2022.
[48] Yaniv Ovadia, Emily Fertig. Jie Ren, Zachary Nado, D. Sculley. Sebastian Nowozin, Joshua Dillon, Balaji Lakshminarayanas and Jasper Snoek. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. In Wallach, H. Larochelle. A. Beygelzimer. F. d'Alché-Bus BBch Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32, Curran Associates, Inc., 2019.
[49] Atahan Ozer, Burak Buldu, Abdullah Akgül, and Gozde Unal. How to combine variational bayesian networks in federated learning. In Workshop on Federated Learning: Recent Advances and New Challenges (in Conjunction with NewrlPS 2022), 2022.
[50] Kaan Ozkara, Navjot Singh, Deepesh Data, and Suhas Diggavi. Quped: Quantized personalization via distillation with applications to federated learning. In A. Beygelzimer, Y. Dauphin, Liang, and J. Wortman Vaughan, editors, Advances in Neurul Information Processing Systems, 2021.
[51] Kaan Ozkara, Antonious M. Girgis, Deepesh Data, and Suhas Diggavi. A statistical framework for person- alized federated learning and estimation: Theory, algorithms, and privacy. In The Eleventh International Conference on Learning Representations, 2023.
[52] Nicolas Papernot, Martín Abadi, Úlfar Erlingsson, Ian J. Goodfellow, and Kunal Talwar. Semi-supervised knowledge transfer for deep learning from private training data. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
[53] Edoardo Patelli, Diego A. Alvarez, Matteo Broggi, and Marco de Angelis. Uncertainty management in multidisciplinary design of critical safety systems. J. Aerosp. Inf. Syst., 12(1):140-1e) 169, 2015.
13

---

[54] Reese Pathak and Martin Wainwright. Fedsplit: an algorithmic framework for fast federated optimization. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 7057-7066. Curran Associates, Inc., 2020.
[55] Felix Sattler, Tim Korjakow, Roman Rischke, and Wojeiech Samek. Fedaux: Leveraging unlaheled auxiliary data federated learning, 2021.
[56] Felix Sattler, Robert Müller, and Wojeiech Samek. Clustered federated learning: Model-agnostee distributed multitask optimization under privacy constraints. IEEE Trans. Neural Networks Learn. Syst. 32(8))3710-337 2021.
[57] Aviv Shamsian, Aviv Navon, Ethan Fetaya, and Gal Chechik. Personalized federated learning using hypernetworks. In International Conference on Machine Learning, pages 9489-9502. PMLR, 2021.
[58] Sidak Pal Singh and Martin Jaggi. Model fusion via optimal transport. Advances in Neural Information Processing Systems, 33, 2020.
[59] Virginia Smith. Khao-Kai Chiang, Maziar Sanjabi, and Ameet S Talwalkar. Federated multi-task learning. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[60] Virginia Smith, Chao Kai Chiang, Maziar Sanjabi, and Ameet S. Talwalkar. Federated multi task learning. In Isabelle Guyon, Ulrike von Luxburg. Samy Bengio, Hanna M. Wallach, Rob Fergus. S. V. N. Vish- wanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 4424 4434, 2017.
[61] Shengyang Sun, Guodong Zhang, Jiaxin Shi, and Roger Grosse. FUNCTIONAL VARIATIONAL BAYESIAN NEURAL NETWORKS. In International Conference Learning Representations, 2019.
[62] Long Tian, Wenchao Chen, Bo Chen, Muyao Wang, Liang Dai, Baolin in Sun, and Mingyuan Zhou. Variational adaptive graph transformer for multivariate time series modeling, 2023.
[63] Ba- Ba-Hien Tran, Simone Rossi, Dimitrios Milios, and Maurizio Filippone. All You Need is a Good Functional Prior for Bayesian Deep Learning. Journal of Machine Learning Research, 23:1-56, 56, 2022.
[64] A. Triastcyn and B. Faltings. Federated learning with bayesian differential privacy. In 2019 IEEE International Conference on Big Data (Big Duta), pages 2587-2596. IEEE Computer Society, 2019.
[65] Elahe Vedadi, Joshua Dillon, Philip Andrew Mansfield. Karan Singhal, Arash Afkanpour, and War- ren Richard Morningstar. Federated variational inference: Towards improved personalization and general- ization, 2023.
[66] Hongyi Wang. Mikhail Yurochkin, Yuckai Sun, Dimitris Papailiopoulos, and Yasaman Khazacni. Federated learning with matched averaging. In International Conference on Learning Representations, 2020.
[67] Kaibin Wang, Qiang He, Feifei Chen, Chunyang Chen, Faliang Huang, Hai Jin, and Yun Yang. Flexifed: Personalized federated learning for edge clients with heterogeneous model architectures. In Proceedings of the ACM Web Conference 2023, WWW '23, page 2979-2999999 Association for Computing Machinery. 2023, ISBN 9781450394161.
[68] Yuxin Wen, Jonas A. Geiping, Liam Fowl, Micah Goldblum, and Tom Goldstein. Fishing for user data in large-batch federated learning via gradient magnification. Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 23668 23684. PMLR, 17-23 Jul 2022.
[69] Andrew G Wilson, Zhiting Hu, Russ R Salakhutdinov, and Eric P Xing. Stochastic variational deep kernel learning. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems. Curran Associates, Inc.
[70] Feng Yin, Zhidi Lin, Yue Xu, Qinglei Kong, Deshi Li, Sergios Theodoridis, and Shuguang. Fedloc: Federated learning framework for data-driven cooperative localization and location data processing. JEEE Open Journal of Signal Processing, 1:187-215, 2020.
[71] Lei Yu. Ling Liu, Calton Pu, Mehmet Emre Gursoy, and Stacey Truex. Differentially private model publishing for deep learning. In 2019 IEEE Symposium on Security and Privacy (SP), pages 332-349, 2019. doi: 10. 201110099......00099..
14

---

[72] Tao Yu, Eugene Bagdasaryan, and Vitaly Shmatikov. Salvaging federated learning by local adaptation. CaRR, abs/2002.04758. 2020.
[73] Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors. Proceedings of the 36th International Conference on Machine Learning. volume 97 of Proceedings of Machine Learning Research, pages 7252-7261, Long Beach, California, USA, 09-15 Jun 2019. PMLR.
[74] Xinwei Zhang, Mingyi Hong, Sairaj Dhople, Wotao Yin, and Yang Liu. Fedpd: A federated learning framework with adaptivity to non-iid data. IEEE Transactions on Signal Processing, 69:6055 6070, 2021.
[75] Xu Zhang. Yinchuan Li, Wenpeng Li, Kaiyang Guo, and Yunfeng Shao. Personalized federated learning via variational Bayesian inference. In Kamalika Chaudhuri, Stefanic Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference Machine Learning, volume 162 Proceedings of Machine Learning Research, pages 26293 26310. PMLR, 2022.
[76] Qingping Zheng, Jiankang Deng, Mingming Gong, Ying Li, and Stefanos 7afeiriou. Global-local bayesian transformer for semantic correspondence, 2023.
15

---

Supplement for "A Bayesian Approach for Personalized Federated Learning in Heterogeneous Settings"
In this supplementary material, first go over the preliminaries of Bayesian learning methods, followed by the pseudo-code of the algorithm used for training our framework. Then, we provide definitions and results used the privacy analysis of the method along with the proof of our privacy budget theorem. We show model calibration metrics and present results demonstrating our method is well-calibrated. We also discuss the details about the alignment dataset, AD, its affect on the performance. include additional experimental results and discuss the communication and computation cost of the procedure.
A Bayesian Learning
Consider learning setting where trying train neural network dataset X' The aim of this setting thus to obtain the set weights, denoted by W, for the corresponding neural network that best fits the data. We could also view a neural network as model that outputs P(yla, W) which is the distribution the label y for a given data point æ under the weights W, for classification this would be the output of the softmax function. Now, the weights of the network can be learnt by Maximum Likelihood Estimation (MLE) for given set of datapoints X
(Zx, y4 M solving the following optimization problem.
W arg max
Σ (yy W)
W
This optimization could be solved by gradient descent based methods and obtains a point estimate of the weight vector, denoted by WALEE
The Bayesian learning methods, on other hand, obtain posterior distribution on the weights given the training data, (P) which as opposed to the point estimates denotes the joint distribution of all the weight parameters of the network over the set of values they are likely to take under the observed data and the prior information encoded in the prior distribution. The predictions for any new data point,
then obtained hy taking expectation of the prediction under the posterior distribution, y E.--p(W xx) [P [ar, w0) Exact inference the posterior distribution, however, is intractable neural networks. Variational inference a traditional approximation method used to obtain an approximation of the posterior weight distribution, and it has also been shown to work for neural networks (29). Specifically, variational inference tries to learn a simpler parameterized distribution ((w(0) from a family of distributions Q by optimizing the parameters such that the new distribution g(w(0") obtained for the optimal value of is close the posterior distribution IP())X)) Precisely, the optimization problem looks like
0" = mig KL[a(W10) [P(]]))]
(1(())) P(W])2)
(6)
((W())
arg min B.((Wl)Ee
mig
(WW10) P(W)
(7)
argmin KI v)] R [og
(8)
e.(WW))cC
where p(W; ob) signifies the prior distribution over weights W parameterized by v. The prior distribution is typically used to encode any previously available information about the weights of the network. The above given objective the same objective as Equation that is used local training in our method.
B Algorithm
The pseudo-code of algorithm used in FedBNN method included the Algorithm [ The Algorithm [ works in the setting when there is a server connected to N clients with each client having local dataset X;; of size ny drawn from the local data distribution D,, and the server has an auxilliary unlahelled dataset called AD. The output of the algorithm is the set of personalized models 4 parameterized by W. for each client i. All W, 's instead of being point estimates, determined by a posterior distribution PP(W---) which is learnt from the data via variational inference. As mentioned in the Section 3.2 the learning procedure first optimizes prior parameters by minimizing Equation aadd then learns the posterior parameters keeping the prior fixed by minimizing Equation S
C Privacy Analysis
Some known results on differential privacy that are used to determine the privacy loss of our algorithm are given in this section and then the proof of the Theorem 4.1 is presented.
16

---

Algorithm FedBNN Algorithm
Input:
number of clients N, number of global communication rounds E, number of local epochs
e, weight vector [W1 W2. ... wNl. noise parameter 7 Output: Personalized BNNs {Øşlé li E [1, parameterized by W; ~ PP(W)))
Server S
AD
for I. to E do
Select subset of clients M₁
for each selected client i E M2 do
(X(X) LocalTraining(,, (X) (t-)),,
end for i ((X)() (1) E.V. wy (X)
end for
Return 21 (E), (2 (E)
LocalTrranning(t (t-1(
Run inference on X to obtain (X)
Generate decorrctedd (X) = (X) [- (1 (X)
for each prior epoch do
Minimize CrossEntropy( (oooeeetdd (X), (X)) to obtain prior parameters yi of the BNN c; end for for each local epoch do
Minimize (L] (Io (ogg,je over (0 <(W110) E Q) to
obtain 0 end for
P(W.) ((V.)0"))
Obtain Monte-carlo samples w!() ; [1, m] from P(W)X)
Compute (X(X) 1 Σラ (x (X:W) w!l))
Return (X)
P(Ww:vix)
p(Diish))
prior
X N
X1
posterior
posterior
Client
Client N
(y(D))
(A
@AAD)
4 (AD)
Ew (,DA
4(AD)
Figure A schematic overview of our method: Local BNN on each client obtain a posterior distribution over local parameters using the prior distribution and local data. These local models generate outputs on the AD using their respective posterior distributions and share these outputs with the server. The server aggregates these outputs and distributes the aggregated output on the AD back to all clients, guiding the prior distribution on each client. These updated prior distributions then further guide the learning the posterior distributions.
17

---

Definition C.1 ((e, 6)- Differential Privacy). A randomized algorithm is (e. 8)-DP if for any two neighboring datasets D and D' that differ in most data point, the output of the algorithm M D and D' is bounded as
P[M(D) E $s & PIM(D)) E S) + 6, AS C Range(MM.
A generalization of differential privacy is known as concentrated differential privacy(CDP). And an alternative form of concentrated differential privacy called zere-concentrated differential Privacy(zCIDP) was proposed enable tighter privacy analysis 12] We will also use the >CDP notion of privacy for our analysis. The relationship between standard DP and zCDP is shown below.
Proposition C.2 ((e, 6D-P and p-2CDP), For a randomized algorithm M to satisfy (e. 8)-DP, it is sufficient for it to satisfy 2² _ ZCDDP And randomized algorithm M that satisfies p-zCDP, also satisfies (é', 8)DDP where
alog
d p+ V 4
As opposed to the notion of DP, the zCDP definition provides tighter bounds for the total privacy loss under compositions, allowing better choice of the noise parameters. The privacy loss under the serial composition and parallel composition incurred under the definition of zCDP was proved by [7]] and recalled below.
Proposition C.3 (Sequential Composition). Consider two randomized mechanisms, Mi and M2, if is P1 D1PCCDP and M2 PCrrCDP, then their sequesntial composition given by (M1(),
is (pr + Pa)-
Proposition C.4 (Parallel Composition), Let mechanism M consists of a sequence of ls adaptive mechanisms, MM, working on randomized partitioo of the D> = (D1, D2.... ........ D&), such that each mechanism Mi is Pt-CDDP and Me IT,11 O2 x Dr + O2 then M(D) =
(Dz))) is
max; Pi-2CDP.
After computing the total privacy loss by algorithm using the tools described above, we determine the variance of the noise parameter a for set privacy budget. The relationship of the noise variance to privacy has been shown in prior works by [9] [1] and is given below.
Definition C.S (L2 Sensitivity). For any two neighboring datasets, D and D' that differ in at most data point, La sensitivity of mechanism M given by maximum change in the La norm the output of M on these two neighboring datasets
(2(M) = sup [IM(D) - M(D')l2 D.D'
Proposition C.6 (Gaussian Mechanism). Consider mechanism M with L2 sensitivity A, if on query q, the output of M is given as M(x) ((z) + o²), then M is A2 202 CCDP
Equipped with the above definitions and results, we now re-state the bound on the privacy loss of algorithm and provide proof below.
Theorem C.7 (Privacy Budget). The proposed algorithm is (c, ddifferrnntl private, if the total privacy budget per glohal communication round per query is set
for E. number of global communication rounds and K number of queries to the algorithm per round.
Proof. After using Gaussian mechanism on each client and adding noise to each coordinate of (AD).. the local mechanism at each client becomes p-ZCDP for p 26z- Since each client outputs the logit representation for each input, i.c., the normalized output the clients, A S 2. The sensitivity, denoted as A, is defined in Definition C.4 in the paper which defines Lz-sensitivity as the maximum change in the La norm of the algorithm'' 's output between two neighboring datasets differing in at most one data point. Let D and D' be two neighboring datasets that differ in one data point present at the 7th row (without loss of generality), and let 4 (D(i.)) :)) be the ne (number of classes) dimensional output probabilities from the model & for the 2t row datapoint in D and (DD((t))) be the output probabilities for the 2th row datapoint in D'. The La sensitivity of 4 is
A²
A A (D) 4(D)ll2 Since all other data-points hetween and D' are identical, the La2 sensitivity of + hecomes-
A ()) = || (DD(i.)
((D'(i,:))l
Now, & (i)).)) and (i)(i))) are both probability distributions, therefore it can be seen that the squared La norm of their difference is bounded by 2, i.e., A (Φ) < 2 (the maximum occurs when & ((((i))) = and (DD(i,)) :)): for two separate indices t: + 0).
18

---

Now, suppose in each global communication round we make K queries each client, then by sequential composition C.3, get EKp, p. for E number of global communication rounds. By parallel composition C.4 the total privacy loss for all clients the maximum of the loss on each client and therefore remains E
Relating it to the (e, 6)-DP from C.2 we get p 4 K
for any 8 >> O.
D Uncertainty Quantification and Calibration
Model calibration is way to determine how well model's predicted probability estimates the model's true likelihood that prediction. Well-calibrated models are much more important when the model decision used critical applications like health, legal etc. because in those cases managing risks and taking calculated actions require confidence guarantee well. Visual tools such as reliability diagrams often used to determine In reliability diagram, model's accuracy the samples is plotted against the
if model is calibrated confidence. perfectly calibrated model results in an identity relationship. Other numerical metrics that could be used to measure model calibration include Expected Calibration Error (ECE) and Maximum Calibration Error (MCE). ECE measures the expected difference between model confidence and model accuracy whercas MCE measures the maximum deviation between the accuracy and the confidence. The definitions and empirical formulas used for caleulating ECE and MCE are as given below.
ECE Ep IP( Y|P -p]-
MCE max |P( YIP = = pl- pE|o.1]
Empirically,
ECE Σ
M
B.
accuracy(B:) (onfidence(B:))
==1
MCE
max accuracy confidence(B:)), .E
where B; is bin with set of indices whose prediction confidence according to the model falls into the range ا_ { F shows the reliahility diagram along with the ECE and MCE scores for method measured
on MNIST and CIFAR-10 dataset the non-IID data setting.
1.0
1.0
0.8
0.8
Accuracy
Accuracy
0.6
0.6
0.4
0.2
0.2
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Confidence
Confidence
(a) Dataset: CIFAR-10, ECE: 0.070, MCE: 0.134 (b) Dataset: MNIST, ECE: 0.032, MCE: 0.156
Figure 3: Reliability diagrams and scores showing model calibration. Figure (a) is for the results corresponding to the CIFAR-10 dataset and Figure (b) for MNIST dataset.
For the next analysis, we consider an approach similar to analysis on uncertainty quantification done in exemplar works in this area 35)48) 48). One of the key requirement of reliable estimates and uncertainty quantification is ensuring high confidence in correct predictions and low confidence in incorrect predictions. To assess whether the proposed method meets this criterion, train method using standard MNIST train dataset, and test dataset well on an -f-ssriibbitt dataset composed NotMNISTIO (featuring
on the MNIST images of alphabets instead of digits). Then, we compute the entropy of the predictive distribution (distribution over output class probabilities) for each dataset and visualize this entropy in the provided Figure ?? In the first corresponding to the in distribution dataset, both Bayesian model and the non- non-Bayesiar model exhibit
19

---

low entropy, expected. However. for the out-of-distribution test dataset, while the non-Bayesian method demonstrates low entropy, method yields high entropy. This observation implies that the non-Bayesian method tends be overly confident predictions on unknown classes, which could pose significant risks practical seenarios, especially in critical applications.
Indistriburian heciieg
Indiittib.ttio Fedloway
203
11.
75
3.-
000
00
Qut-ef-ittiiutton roddug
Oft-of-dsttBill
03
00
Figure 4: Distribution of the entropy of class-probability distributions across different clients demon- strating the confidence of methods in predicting on in-distribution vs out-of-distribttion data.
E
Alignment Dataset (AD)
10-
60
30
4020
5000
Figure 5: Ablation study comparing the affect of AD size the performance. The included results are for CIFAR-I0 dataset in the small data setting with non-lID partitions and heterogeneous clients.
In FedBNN, the alignment dataset (AD) is used to achieve collaboration across clients. Since the only assumption AD is for it to be of the same domain the target application, there is practical constraint obtaining the AD real-world settings. In many cases could obtained from web, for example images from common datasets in Huggingface, texts from Wikipedia, Reddit etc. Furthermore, the server having its own dataset in addition to the private data on the clients is common in various cross-silo and cross-device applications. For instance, hospitals with access to patient records may combine this data with private patient data collected from individual devices such wearables sensors for federated learning, might leverage open-source code along with private code repositories from developers to enhance generative models, delf-driving car companies may collect their own data and utilize it alongside private data collected from customers' vehicles, and many more. The use of AD is not different from how several other methods use an additional dataset for augmentation, and several other methods have used labelled or unlabelled auxiliary
code generation applications
20

---

Table 3: Effect of varying distribution of AD on the clients' performance for the non- secting with CIFAR-10 10 dataset and 20 clients where each client has data for the 5 different classes.
Architecture Setting
Local Training CIFARI0(10) CIFARI0(8) CIFARIO(S) CIFARI0(2)
SVHN
Homogeneous Architectures
64.3.1 0.36
72.7 = 0.15
69.7 = 0.28
68.8
0.97
67.2
1.5
70.1 0.18
Heterogeneous Architectures
61.2+ 0.17
71.6 + 0.93
68.4 + 0.80
68.8 + 1.4
68.1 + 1.9
69.3 + 0.8
dataset to improve the performance and privacy of the FL algorithms [2 36 [5S]. The effect of size of AD on the performance of models is demonstrated in Figure [S For CIFAR-I0 dataset in the small data and non-IID setting. In that figure, we observe that when the size AD small the performance of the model is low but as the size of AD increases the performance increases up to a point and becomes constant afterwards. The number of data points in AD that are required to achieve good improvement in the model performance is small and practical.
We also vary the distribution of the AD being used and test the final performance of the models and report it in Table [] We these experiments on 20 clients for CIFAR-10 dataset where each client had access only 5 of the 10 classes and each client helonged to the medium data setting. For the first experiment, we use a held-out dataset from the CIFAR-10 10 data as AD but vary the composition of the dataset by changing the distribution of the classes present in the AD. for example. CIFARIO(1O) is composed of all 10 classes present in the CIFAR- 10 dataset but CIFARI0(2) is composed of only 2 out of the 10 classes present in the AD and likewise. We also test the performance of our method when significantly different dataset SVHN consisting of the colored house number images is used. T suggests that the performance of the method even with different datasets as AD always improves and that the gain between local training and the proposed procedure is better highlighted in the heterogeneous architecture settings, since there local client capacitics and model architectures differ significantly and clients are able to utilize the peer knowledge to learn better models locally. We observed that even for different and dissimilar data distributions in AD, it is possible to obtain a value for the parameter such that the final performance of the local client model with collaboration is better than the model independently trained locally on the client. The best results for CIFAR-10 10 classification are seen when AD is composed of held-out set from all 10 classes of CIFAR- 10 denoted CIFAR 10) which is expected. Then, as the composition of AD is changed from 10 classes to random 8, 5 and classes of CIFAR-10 10 (denoted as CIPARIO(8) CIFAR 10(5) and CIFAR10(2) respectively) the performance keeps on decreasing. We see that the performance of the same task with SVHN as AD is only strictly better than CIFAR 10(2) as AD. We believe that the SVHN dataset as AD works hetter than the CIFAR 10(2) hecause it provides more variability in the data distribution. A similar observation for knowledge distillation, which noted that the out-of-domain unlabelled data for distillation can perform even better. Moreover, the parameter 7 controls the amount of global knowledge to be incorporated on each client and with appropriately set 7, the AD also helps achieve regularization for the local client model such that the local models do not overfit to the relatively smaller local datascts and generalize better.
also recorded in FedAUX (55) which also uses additional unlabelled data
F Communication and Computation Efficiency
Communication Cost In FedBNN, each global communication round requires that the server sends the alignment dataset to all the clients and the clients upload the outputs of their respective models on the common dataset AD. Since AD is publicly available dataset, AD could be transmitted to the clients by specifying the source and the indices, and does not really needs to he communicated across the channel. The client output on AD, on the other hand, depends on the number of instances in AD, let's call it K, therefore, the total communication cost in each round of our method is O(K). As shown in Figure having K 2000 gives a good performance. The communication cost between the clients and the server, thus, is also invariant of the number of model parameters which tends to run in millions. This allows our method to be much more communication efficient as compared to the conventional FL algorithms and other Bayesian FL methods that transmit model parameters or parameter distributions in each communication round, making it practically more useful.
Computation Cost Similarly, the computation cost a FL procedure involves the costs incurred in local training at the individual clients and the cost of aggregation at the server, both of which are discussed below.
Server-side computation cost The server side computation cost arises from the need to aggregate knowledge obtained from individual clients. In the state-of-the-art bayesian FL algorithms, the server aggregates posterior distributions for each weight parameter in the neural network obtained from various clients. The number of such weight parameters typically run in millions. In our method we do not aggregate the parameter distributions but achieve collaboration by aggregating the client outputs the AD (with size 2000). thus the server side computation cost in our method is many orders
21

---

Table 4: Performance comparison as function of the privacy guarantee.
Privacy (c) per round
Test Accuracy
75.5 %
R3 0.1
71.3 %
R3 0.01
68.6 %
0.001
62.2 %
0.0001
59.6 %
Table 5: Test aceuracy comparison with more number of clients (500) in the setting.
Method
Test Accuracy
PFedGP
53.2 + 0.4
PFedBayes
52.9 = 0.8
Ours(Homo)
56.1 1 0.3
Ours(Hctcro)
54.7 + 1.0
of magnitude lower than the conventional methods and does depend on the number of model parameters. This makes our method much more efficient and scalable than existing federated bayesian solutions.
Client-side computation cost The client-side computation cost is mostly determined by the cost of training Bayesian Neural Network the client side, which turn depends on type of inference procedure used for obtaining the posterior distribution over weights for each parameter of the neural network. In the proposed work, the method used for inference is Bayes by Backprop which uses the gradient computations similar to backpropagation(BP) algorithm butions where the posterior distrihutions are characterized by the mean and standard deviation. re- re-parameterization trick used to compute the mean and std of the distributions from the back- propagated gradients. Thus the cost for obtaining the posterior distributions similar to the cost of backpropagation. Moreover, since the method only uses gradient updates, the optimizations used for SGD like asynchronous SGD could be readily used for obtaining the posteriors. An unrelated but similar algorithm (28) does probabilistic backpropagation
obtain the posterior distri-
train BNNs and shows that the average run time of probabilistic BP is not higher than that of BP.
To summarize, the communication cost and the server-side side computation cost of the proposed method is orders of magnitude lower than that of the other Bayesian baseline methods. On the other hand, the client-side computation cost is determined by inference procedure used to obtain posterior distributions and which Bayes by Backprop provides an efficient mechanism. Several works the recent past have discussed the use related Bayesian inference based methods for training uncertainty-awar transformers [76 62) 42) proving that Bayesian methods not limited to use in simpler models. And therefore, framework can also be extended to apply in settings where much larger neural networks required.
G Additional Experiments
Since the amount of noise required be added to the client's outputs via the Gaussian Mechanism directly proportional to the guaranteed privacy, we the affect of the privacy guarantee on the performance the proposed framework by comparing the performance of method with varying e and The results are reported in Table 4 We observe that, as expected, when we reduce the amount of
Privacy vs Performance
5 = 10 4 privacy loss in each iteration by adding more noise to the clients' outputs going to the server, the performance of the method drops. However the drop in performance in all the cases is drastic the clients tune the level personalization or global knowledge required by appropriately setting the parameter in Equation 3
More clients To test the performance of the proposed method when a large number of clients are involved the setup, we did additional experiments with 500 clients and non-lIl setting with classes per client in the medium data setting on the CIFAR-10 dataset where in cach communication round only 10% of the clients are selected for participation and 0.7. The obtained results at the end of 200 communication round are reported in Table 5 We observe that the homogeneous version of method is better than the baselines by significant margin and the heterogeneous version is slightly better than the baselines.
22

---

NeurIPS Paper Checklist
1. Claims
Question: the main claims made the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We support all the claims made in the abstract and introduction by thorough experimental evaluation in the paper. Please refer Section for details.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much results can he expected to generalize to other settings.
It is fine to include aspirational goals as motivation long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
Please refer Section d
Guidelines:
The answer NA mcans that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well- speciflcciiii asymptotic approximations only holding locally). The authors should reflect how these assumptions might he violated practice what the implications would be.
The authors should reflect on scope the claims made, e.g., the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't It acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play important role in developing norms that preserve the integrity the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Question:
Answer:
[Yes]
The detailed analysis is presented in the Section 4 and Appendix C due to space constraints.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and reffeeeffeee
23

---

All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly refereneed.
Experimental Result Reproducibility
Question: Does the paper fully disclosc all the information needed to reproduce the main experimental results of the paper the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Please refer Section B and subsection 5.1 the experimental details.
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible important, regardless whether the code and data are provided or not.
the contribution is dataset and/or model, authors should describe steps taken make their results reproducible or verifiable.
Depending on contribution, reproducibility can be accomplished various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access the model. general. releasing code and data often one good way to accomplish this, reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (c.g., in the case of a large language model), relcasing of model checkpoint, or other means that appropriate to the rescarch performed.
While NeurlPS does not require releasing code, the conference does require all submissions provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily reproduce that algorithm.
algorithm, the paper should make clear how
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g.. a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (c.g.. with open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source source models, it may be that access to the model is limited in some way (e.g.. registered users). but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does paper provide open access to the data code. with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: The code will be released on acceptance of the paper.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurTPS code and data submission guidelines (https://nips... http:a::/niip..ccppbbbll/ cc/public/ http//ooddbiiiiiiooooooo
for more details.
While we encourage release of code and data, we understand that might not possible,
so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source henchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: : :/nniss//////////////e:: cc/public/ qul.dess/CooeSuubmiiissooooooo.. for more details.
24

---

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) recommenedd, but including URLs to data and code is permitted.
6. Experimental Setting/Detann Details
Question: Does the paper specify all the training and test details (c.g,, data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Please refer Section 5 and subsection 5.1 for the experimental details.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate informa- tion about the statistical significance of the experiments?
Answer: [Yes]
Justification: Please refer Section and subsection 5. Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should given (e-g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error than state that they not verified.
CL, if the hypothesis of Normality of errors
asymmetric distributions, authors should be careful not show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). crror bars are reported tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question: For each experiment, docs the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification:
Please refer Section and subsection 5.1 for these details.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
25

---

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make into the paper).
9. Code Of Ethies
Question: Does the research conducted the paper conform. every respect, with NeurIPS Code Ethics hhttts://neur/// http:///errrpps.ccpppiii///////a httpp//aaaccetetaaaeuuuaaaaaaaaa
Answer:
[Yes]
Justification:
We understand the NeurlPS Code of Ethics and the work conforms to that.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e-g., if there is a special consideration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [NA]
Justification: The work in this paper does not have societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no socictal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deploy- of technologies that could make decisions that unfairly impact specific groups), privacy considcrations, and security considerations.
The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it not needed point out that generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated relcase of models, providing defenses in addition to attacks, mechanisms for monitor- ing misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible relcase of data or models that have a high risk for misuse (e.g. pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: NA.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors take this into account and make best faith effort.
26

---

12. Licenses for existing assets
Question: Are creators or original owners assets (e.g., code, data, models), used the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected"'
Answer: [Yes]
Justification: We have carefully cited the datasets and codes used in the paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, ppaperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license a dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
this information not available online, the authors are encouraged reach out asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: NA.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their sub- missions via structured templates. This includes details about training. license, limitations.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, applicable, as as details about compensation (if any)?
Answer: [NA]
Justification: NA.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: NA.
Guidelines:
27

---

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this the paper.
recognize that the procedures for this may vary significantly between institutions and Iocations, and we expect authors adhere to the NeurIPS Code of Ethics the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applica- ble), such as the institution conducting the review,
28


Reviews:
[{"review_id": "yzd0PrKTRK", "reviewer": "Reviewer_rtXt", "summary": "The paper proposes FedBNN, a Bayesian approach for personalized federated learning. The approach relies on the availability of an auxiliary small public unlabelled dataset, called Alignment Dataset, that can be used as a mean to distill knowledge across clients. In FedBNN, clients maintain an estimate of a  posterior distribution over the model parameters, which are updated locally using Bayes-by-Backprop. The clients' posteriors are aggregated through Monte Carlo sampling on the  Alignment Dataset. FedBNN has a differentially private implementation which consists in adding a Gaussian noise to the output of the forward pass of the local models on the Alignment Dataset. \n\nThe paper conducts thorough numerical experiments to quantify the performance of the FedBNN, and demonstrates that it outperforms SOTA baselines on a wide range of datasets and heterogeneity levels and types.\n\n----\nPost Rebuttal\n\nThe rebuttal addresses my concerns on the DP guarantees, as such I raise my score to 6.", "questions": "* Could elaborate more on the DP analysis. In particular, I would appreciate if you can formally prove why $\\Delta < 2$ in  the proof of Theorem 4.1 in Appendix C. Also, could you please explicitly show the effect of $\\sigma_g$ on the DP guarantees in Theorem 4.1.", "limitations": "FedBNN relies on the availability of a public auxiliary dataset.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* The paper is overall well written and easy to follow.\n* The proposed FedBNN handles both statistical and system heterogeneity. In particular, it allows each client to have a personalized architecture. \n* The numerical experiments are rigorous and considers a relatively large number of settings and competitors.", "weaknesses": "* FedBNN relies on the availability of a public auxiliary dataset. It is true that an auxiliary dataset is often available, but sometimes it is none. \n* I am not sure how to interpret Theorem 4.1 and if the DP mechanism of the paper is correct. On one hand, Theorem 4.1 does not take into consideration the noise multiplier of the DP mechanism, which should be related to $\\sigma_g^2$. On the other hand, the DP mechanism presented in the paper does not have any kind of clipping or normalization, which is usually expected. The proof in Appendix C, claims that $\\Delta$ is smaller than two, but it only provides a short justification. (*Note: I am not a DP expert*).\n\n**Given the current doubt on Theorem 4.1, I am leaning towards rejection, but I am willing to adjust my rating when this point is clarified.**", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer rtXt,\n\n\nWe are thankful to the reviewer for the thorough feedback on our work. \n\nIt is encouraging for us to see that the reviewer found the paper easy to read and follow, appreciated the rigorous \"*experimental evaluation that considers a wide range of settings and competitors*\", and recognized the method's effectiveness in addressing both statistical and system heterogeneity.\n\n\nBelow we address the reviewer's concerns on the DP analysis\t-\n\n-- *Dependence on $\\sigma$*\n\nWe acknowledge the doubt regarding not including the noise-parameter $\\sigma$ in the main theorem, Theorem 4.1, of the paper. However, we elaborate on this in the proof of the theorem present in the Appendix C (Theorem C.7). The first line in the proof suggests that $\\rho = \\dfrac{\\Delta^2}{2 \\sigma^2}$ establishing the relationship between the noise-parameter $\\sigma$ and the $(\\epsilon, \\delta)$ parameters of differential privacy. We have also established this relationship in the Theorem 4.1 in the paper itself in the revised version of the paper.\n\n\n\n-- *Clarification on the bound on $\\Delta$*\n\nThe sensitivity, denoted as $\\Delta$, is defined in Definition C.4 in the paper which defines $L_2$-sensitivity as the maximum change in the $L_2$ norm of the algorithm's output between two neighboring datasets differing in at most one data point. Let $D$ and $D'$ be two neighboring datasets that differ in one data point present at the $i^{th}$ row (without loss of generality), and let $\\Phi(D(i,:))$ be the $n_c$ (number of classes) dimensional output probabilities from the model $\\Phi$ for the $i^{th}$ row datapoint in $D$ and $\\Phi(D'(i,:))$ be the output probabilities for the $i^{th}$ row datapoint in $D'$. The $L_2$ sensitivity of $\\Phi$ is -\n\n$\\Delta (\\Phi) = || \\Phi(D) - \\Phi(D')||_2$\n\nSince all other data-points between $D$ and $D'$ are identical, the $L_2$ sensitivity of $\\Phi$ becomes -\n\n$\\Delta (\\Phi) = || \\Phi(D(i,:)) - \\Phi(D'(i,:))||_2$\n\nNow, $\\Phi(D(i,:))$ and $\\Phi(D'(i,:))$ are both probability distributions, therefore it can be seen that the squared $L_2$ norm of their difference is bounded by 2, i.e., $\\Delta(\\Phi)^2 \\leq 2$ (the maximum occurs when $\\Phi(D(i,:))_k = 1$ and $\\Phi(D(i,:))_l = 1$ for two separate indices $k \\neq l$).\n\nWe greatly appreciate the opportunity provided by the reviewer to provide better understanding of the privacy analysis. We have now included these details in the revised version of the paper.\n\nBest,\n\nAuthors", "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": "", "extra_scores": {}}, {"review_id": "ncXx8IcdGM", "reviewer": "Reviewer_yjRB", "summary": "This work integrates Bayesian neural networks, knowledge distillation, and differential privacy into the federated learning framework. The resulting framework can maintain privacy during training and provide uncertainty estimates for its predictions. Empirical results demonstrate that the proposed method outperforms previous works in the given experimental setting.", "questions": "Since only features of the public dataset are transferred, this framework should be communication efficient. The authors could consider plotting accuracy versus transmitted bits to emphasize the strength of the proposed method.", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. This work is well-written. The authors have clearly explained the motivation, methodology, and experimental settings.\n2. This work addresses data and system heterogeneity and privacy, touching on all the major challenges of federated learning.", "weaknesses": "1. Bayesian frameworks and knowledge distillation have been independently integrated into federated learning and studied by many works. Although this work combines them with differential privacy effectively and shows good results, it has limited novelty as an academic research paper.\n2. There are limited experimental results. The main results are presented in Table 1, where the experimental setting considers 20 clients and label distribution shift. Given the complexity of practical situations, the authors could consider other experimental settings to demonstrate the generality of the method, such as different types of data heterogeneity, varying numbers of clients, and different numbers of data points across clients. I note that in the appendix, the authors also provide results for 500 clients, but there are only a few numbers. A more comprehensive study similar to Table 1 would be preferable.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s valuable feedback. We are pleased to see that the reviewer finds the work well-written, noting that \"*The authors have clearly explained the motivation, methodology, and experimental settings*\" and recognizing that \"*This work addresses data and system heterogeneity and privacy, touching on all the major challenges of federated learning\".\n\nWe address reviewer's key concerns below -\n\n-- *\"Bayesian frameworks and knowledge distillation have been independently integrated into federated learning and studied by many works. \"*\n\nWe agree with the reviewer that the bayesian framework and the knowledge distillation have been independently integrated into FL by previous works. However, our contributions are particularly significant for the following reasons -\n\n1.  *Knowledge distillation in Bayesian learning itself is non-trivial*. While knowledge distillation is a well-established method, the transfer of knowledge between Bayesian models has remained inadequately addressed. Our work introduces a novel method that facilitates collaboration across client-specific Bayesian models by transferring knowledge through a prior specification mechanism in the output space. This contribution significantly advances the field of Bayesian knowledge distillation, in both centralized as well as federated learning settings. \n2. *Bayesian learning across heterogeneous clients in FL is novel.* The designed collaboration mechanism, is pivotal in enabling efficient collaboration across models of heterogeneous architectures (encompassing varying size and shape), significantly advancing the FL and Bayesian-FL procedures for heterogeneous settings. By means of such a procedure, the small compute capacity clients are able to gain about 10 increase in the performance through collaboration with large compute capacity clients (as shown in Table 2 and Figure 1 in the paper). This is in contrast to the conventional Bayesian-FL procedures that work with homogeneous models and would either leave out the small compute clients out of the FL process or reduce the model capacity on the large compute clients leading to under-performance. \n3. *Incorporating a privacy-preserving Bayesian FL mechanism is novel*. As the reviewer correctly points out our method also incorporates privacy along with Bayesian learning in heterogeneous FL settings, making it more usable for a diverse set of critical applications.\n\n-- *\"Experimental Results\"*\n\nWe would like to clarify to the reviewer that apart from the results in Table 1 and Appendix, we also present results across many different settings in Figure 1, where performance is compared across varying degree of heterogeneity in the non-IID and IID-seeting. For example, like the reviewer said, performance with different number of data points across clients is depicted in Figure 1(c) of the paper, and the varying compute heterogeneity is depicted in Figure 1(a) and  Figure 1(b). \n\nThanks to the reviewer for this suggestion, we include more experiments on varying degree of heterogeneity below. Given the limited time for author response, we compare the performance of our method under varying degrees of heterogeneity against key Bayesian and the most competitive Bayesian baselines identified from previous experiments. To evaluate performance under different degrees of heterogeneity, we utilize the Dirichlet distribution. Specifically, for each class in the dataset, we distribute the data among 20 clients according to a Dirichlet distribution with parameter $\\alpha$, smaller values of  $\\alpha$ create more heterogeneous distributions, while larger values create more homogeneous distributions. We conduct this experiment on the CIFAR-10 dataset with a 20-client setting, and report the results in the table below. Note that when  $\\alpha$ is low, the number of classes per client is less than 10, leading to an easier task and higher accuracy for all methods compared to the setting when  $\\alpha$ is high, and all 10 classes are present on each client, resulting in a more complex task. We have also added these new results to the revised version of the paper.\n\n| Method                 | $\\alpha$ = 0.1 | $\\alpha$ = 1   | $\\alpha$ = 10  |\n| ---------------------- | -------------- | -------------- | -------------- |\n| Local Training         | 70.1 $\\pm$ 1.5 | 54.5 $\\pm$ 1.8 | 47.3 $\\pm$ 0.7 |\n| FedAvg (non-Bayesian)  | 66.2 $\\pm$ 1.4 | 53.7 $\\pm$ 1.2 | 49.0 $\\pm$ 0.5 |\n| FedProx (non-Bayesian) | 66.9 $\\pm$ 1.1 | 56.8 $\\pm$ 0.9 | 50.1 $\\pm$ 0.9 |\n| pFedGP (Bayesian)      | 65.7 $\\pm$ 1.2 | 58.3 $\\pm$ 1.9 | 51.8 $\\pm$ 1.1 |\n| pFedBayes (Bayesian)   | 69.3 $\\pm$ 1.4 | 60.1 $\\pm$ 1.7 | 52.2 $\\pm$ 1.2 |\n| FedBNN (Ours)          | 72.3 $\\pm$ 1.9 | 62.5 $\\pm$ 1.4 | 54.3 $\\pm$ 0.7 |", "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": "", "extra_scores": {}}, {"review_id": "OO6QffWmg7", "reviewer": "Reviewer_4fxY", "summary": "The authors propose FedBNN, a novel personalized federated learning (FL) framework that leverages Bayesian principles to address challenges posed by heterogeneous data and computational resources among clients. FedBNN uses Bayesian neural networks to enable robust local training on small datasets by quantifying uncertainties. A novel collaboration mechanism involves sharing priors in the functional space of the networks, rather than directly sharing model parameters, to accommodate clients with varying computational capacities. The approach also includes a differentially private version with formal privacy guarantees. Experiments on standard FL datasets demonstrate that FedBNN outperforms existing methods, particularly in heterogeneous settings and under strict privacy constraints. The main contributions include improved robustness and efficiency in personalized FL, a novel collaboration method using functional space priors, and a formal differential privacy guarantee applicable to general settings.", "questions": "See Weaknesses above.", "limitations": "NA", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents a comprehensive evaluation of their FedBNN framework across multiple dimensions. The authors demonstrate its effectiveness on three major datasets (MNIST, CIFAR-10, CIFAR-100) under various heterogeneous settings, comparing against both Bayesian and non-Bayesian baselines. They provide detailed ablation studies to justify their design choices, including the use of functional space priors and the auxiliary dataset for collaboration. The experiments thoroughly explore the method's performance under different types of heterogeneity (data resources, compute resources, and statistical distribution), showcasing its robustness and adaptability. The authors also present a formal privacy analysis, demonstrating the framework's applicability in privacy-sensitive scenarios. Overall, the extensive experimental evaluation, coupled with the theoretical foundations, provides strong evidence for the effectiveness and significance of their proposed approach in addressing real-world challenges in federated learning.", "weaknesses": "While the authors present an innovative Bayesian approach to personalized federated learning (FL), several improvements are needed. The novelty is somewhat diminished by existing methods like FedPop [1]; clearer differentiation and more comprehensive comparisons are necessary. The experimental evaluation is robust but lacks diversity in dataset types and real-world applications; incorporating more varied and complex datasets would better demonstrate generalizability. Scalability concerns are not adequately addressed, particularly regarding the computational overhead of Bayesian neural networks. The privacy analysis should delve deeper into the trade-offs between noise levels and model accuracy. More thorough ablation studies are needed to isolate the impact of individual framework components, such as functional space priors versus traditional weight space priors.\n[1] Kotelevskii, N., Vono, M., Durmus, A., & Moulines, E. (2022). Fedpop: A bayesian approach for personalised federated learning. Advances in Neural Information Processing Systems, 35, 8687-8701.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s detailed feedback. We are encouraged to see that the reviewer acknowledges the strength of our work, stating that \"*Overall, the extensive experimental evaluation, coupled with the theoretical foundations, provides strong evidence for the effectiveness and significance of their proposed approach in addressing real-world challenges in federated learning*\".\n\nWe address reviewer's key concerns below -\n\n-- *Differentitaion with FedPop*\n\nWhile a capability wise contrast of our method with the existing methods including FedPop is included in the Table 2 in Appendix, here we provide a detailed comparison of our method with the reference suggested by the reviewer, named FedPop.\n\n1. *FedPop can only work in homogeneous FL settings as opposed to our method that can work in heterogeneous settings too*. The FedPop method in [1] assumes a hierarchical statistical model for personalised FL across clients. The hierarchical model constitutes of an unknown population parameter also referred to as prior (denoted by β) from which the individual local parameters for each client are sampled (denoted by z1,z2,..) where the variance of the population parameter determines the heterogeneity of the client specific parameters. Since the client's local parameters are sampled from the same distribution (also called prior), they have identical shape and form and therefore these methods work only in settings when the clients architectures are homogeneous. This is also evident in the aggregation mechanisms used in the algorithms  where they perform element wise aggregation over the parameters. Since the client's local parameters are sampled from the same distribution (the prior), they have an identical shape and form. As a result, these methods are effective only in settings where the clients' architectures are homogeneous. This is also evident in the aggregation mechanisms used in the algorithms, which perform element-wise aggregation over the parameters. However, the problem setting we address involves Bayesian learning in a heterogeneous FL setting, where clients have different computational and data resources. In such a scenario, it is not feasible to utilize FedPop.\n2. *Our method is much more communication efficient*. Since the FedPop method involves transmitting model parameters between clients and the server in each communication round, its communication cost is proportional to the number of model parameters, which can run into millions. In contrast, our method only transmits the outputs on the alignment dataset between clients and the server in each communication round. This significantly reduces the communication cost, as it involves outputs for approximately 5000 data points, which is much smaller in size.\n3. *Incorporating a privacy-preserving Bayesian FL mechanism is novel* Our work presents a privacy-preserving method for Bayesian FL along with a bound on privacy loss, no privacy analysis or privacy-preserving algorithm is included in FedPop.\n\n-- *Privacy Analysis*\n\nWe clarify to the author that the trade-off between model accuracy and privacy guarantee is included in Table 4 in the Appendix.  Since, there is a direct relationship between the noise parameter and the privacy budget as established in Theorem 4.1, the results in Table 4 are also applicable for the trade-off between model accuracy and noise parameter.\n\n-- *Functional space priors versus traditional weight space priors*\n\nWe appreciate the reviewer's suggestion for more thorough ablation studies to isolate the impact of individual framework components, such as functional space priors versus traditional weight space priors.  Knowledge transfer between Bayesian models through priors in functional space is one of the key contributions of this work. These functional space priors are essential to our method, enabling effective knowledge transfer and collaboration across client-specific models in heterogeneous federated learning settings.\n\nTo identify the effect of functional space priors versus weight space priors, we can compare our method using functional space priors against a baseline with traditional weight space priors. We direct the reviewer to the results in Table 1, which compares the pFedBayes method (which uses the weight soace priors) and our method in the homogeneous setting. Additionally, the results in Figures 1(a) and 1(b) can be used to compare the impact of functional space and weight space priors across varying degrees of heterogeneity among clients. These results could be used to demonstrate the impact of functional space priors in achieving superior performance.", "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": "", "extra_scores": {}}, {"review_id": "T3WJPXS6Wm", "reviewer": "Reviewer_PtiT", "summary": "The authors designed a federated learning algorithm that leverages a Bayesian network to enhance the model's robustness, and uses data distillation methods to reduce overhead and adapt to scenarios with different client model structures. The authors theoretically discussed the differential privacy characteristics of the method and demonstrated its effectiveness through experiments.", "questions": "1. Please introduce the difficulties and improvements in integrating the Bayesian network into the data distillation-based federated learning framework.\n2. What is the performance of this method under other data heterogeneity or model structures?", "limitations": "The authors discussed possible social impact and application and also touched on the potential limitations of the method.", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The paper formally discusses the characteristics of differential privacy, ensuring privacy protection.\n2. Experiments about distillation dataset sizes are quite insightful.", "weaknesses": "1. The authors stated that the paper has two contributions: 1) to address the issue of small data volumes per client by using a Bayesian network; 2) to solve the overhead problem brought by the Bayesian network and the model heterogeneity issue by using a distillation method. However, these two methods seem relatively independent and both have their own rich related works. It appears the authors simply combined these two works directly, resembling an A+B work, which is less compelling.\n2. The authors used a distillation scheme, but the distillation scheme requires the server to maintain public datasets. However, in many scenarios, this is very difficult to achieve, which is a major criticism of similar distillation methods.\n3. The experiments are insufficient. As a paper that mainly measures performance through experimental results, the experimental part of the main text is only 2.5 pages (of which 1.5 pages are explanations of the experimental setup). Although the authors added some additional experiments in the appendix, it is still far from sufficient, and the overall experiment was completed under only one setting. It is recommended that the authors at least consider adding the following experiments: 1) Since the paper involves data heterogeneity, different data heterogeneity scenarios (such as Dirichlet non-IID) should be tested; 2) Expand to more common neural network structures instead of using the authors' self-constructed simple CNNs. Even if the client's computing power is insufficient, lightweight networks such as MobileNet should be considered.\n4. From the perspective of paper presentation, this paper has considerable room for improvement, such as: 1) The paper does not provide a framework diagram of the method, making it very difficult to understand the author's method; 2) The writing style is not concise and clear enough, such as the second point in the contributions; 3) The margin between formulas and text, the width of tables, are not well adjusted.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate reviewer's insightful comments and suggestions. We are encouraged by reviewer's positive remarks about the formal privacy guarantee and the experiments on the alignment dataset.\n\nWe address reviewer's specific concerns below -\n\n-- *Clarification on the contribution*\n\nWe agree with the reviewer that the bayesian framework and the knowledge distillation have been independently integrated into FL by previous works. However, our contributions are particularly significant for the following reasons -\n\n1.  *Knowledge distillation in Bayesian learning itself is non-trivial*. While knowledge distillation is a well-established method, the transfer of knowledge between Bayesian models has remained inadequately addressed. Our work introduces a novel method that facilitates collaboration across client-specific Bayesian models by transferring knowledge through a prior specification mechanism in the output space. This contribution significantly advances the field of Bayesian knowledge distillation, in both centralized as well as federated learning settings. \n2. *Bayesian learning across heterogeneous clients in FL is novel.* The designed collaboration mechanism, is pivotal in enabling efficient collaboration across models of heterogeneous architectures (encompassing varying size and shape), significantly advancing the FL and Bayesian-FL procedures for heterogeneous settings. By means of such a procedure, the small compute capacity clients are able to gain about 10 increase in the performance through collaboration with large compute capacity clients (as shown in Table 2 and Figure 1 in the paper). This is in contrast to the conventional Bayesian-FL procedures that work with homogeneous models and would either leave out the small compute clients out of the FL process or reduce the model capacity on the large compute clients leading to under-performance. \n3. *Incorporating a privacy-preserving Bayesian FL mechanism is novel*. As the reviewer correctly points out our method also incorporates privacy along with Bayesian learning in heterogeneous FL settings, making it more usable for a diverse set of critical applications.\n\n-- *On model structures in Experiments*\n\nWe would like to kindly clarify to the reviewer that we did not only use \"self-constructed simple CNNs\" for the experiments, as mentioned in the review. In fact, we used a combination of popular VGG-based models and CNNs to demonstrate the heterogeneity in our setting. This approach highlights the stability and performance of our method under such drastic heterogeneity. Also, under the homogeneous client setting in the experiments (detailed in subsection 5.1 of the paper) all the clients train VGG based models, which aligns with the reviewer's comment on using \"more common neural network structures\".\n\n-- *On different data heterogeneity in Experiments*\n\nThanks to the reviewer for this suggestion, we include more experiments on varying degree of heterogeneity below. Given the limited time for author response, we compare the performance of our method under varying degrees of heterogeneity against key Bayesian and the most competitive Bayesian baselines identified from previous experiments. To evaluate performance under different degrees of heterogeneity, we utilize the Dirichlet distribution. Specifically, for each class in the dataset, we distribute the data among 20 clients according to a Dirichlet distribution with parameter $\\alpha$, smaller values of  $\\alpha$ create more heterogeneous distributions, while larger values create more homogeneous distributions. We conduct this experiment on the CIFAR-10 dataset with a 20-client setting, and report the results in the table below. Note that when  $\\alpha$ is low, the number of classes per client is less than 10, leading to an easier task and higher accuracy for all methods compared to the setting when  $\\alpha$ is high, and all 10 classes are present on each client, resulting in a more complex task. We have also added these new results to the revised version of the paper.\n\n| Method                 | $\\alpha$ = 0.1 | $\\alpha$ = 1   | $\\alpha$ = 10  |\n| ---------------------- | -------------- | -------------- | -------------- |\n| Local Training         | 70.1 $\\pm$ 1.5 | 54.5 $\\pm$ 1.8 | 47.3 $\\pm$ 0.7 |\n| FedAvg (non-Bayesian)  | 66.2 $\\pm$ 1.4 | 53.7 $\\pm$ 1.2 | 49.0 $\\pm$ 0.5 |\n| FedProx (non-Bayesian) | 66.9 $\\pm$ 1.1 | 56.8 $\\pm$ 0.9 | 50.1 $\\pm$ 0.9 |\n| pFedGP (Bayesian)      | 65.7 $\\pm$ 1.2 | 58.3 $\\pm$ 1.9 | 51.8 $\\pm$ 1.1 |\n| pFedBayes (Bayesian)   | 69.3 $\\pm$ 1.4 | 60.1 $\\pm$ 1.7 | 52.2 $\\pm$ 1.2 |\n| FedBNN (Ours)          | 72.3 $\\pm$ 1.9 | 62.5 $\\pm$ 1.4 | 54.3 $\\pm$ 0.7 |\n\n\n\n-- *Presentation Issues*\n\nWe sincerely thank the reviewer for these valuable suggestions. We have included a new figure providing an overview of our method in the revised version of the paper. For the reviewer's reference, the figure and its caption are available at this link: https://i.imgur.com/ZBC2sMu.png.\n\nWe also apologize for the discrepancies between the margin and text in certain sections. These adjustments were made to fit the content within the 9-page limit. With the additional page allowed for the camera-ready version, these issues will be easily resolved.", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper was a borderline case, but the clarifications and additional experiments provided in the rebuttal helped to lift some concerns, leading to a couple of reviewers increasing their score slightly. Overall, this work introduces a novel and interesting Bayesian approach to personalized FL. I recommend acceptance.

Important note: the authors have committed to add further experiments, as requested by some reviewers. I therefore ask the authors to incorporate these changes in the final version.
Decision: Accept (poster)