A Bayesian Approach for Personalized Federated Learning in Heterogeneous Settings

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 05:16:16

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.

Main paper: 
 introduction: 
 The vast majority of research on Federated Learning (FL) takes an optimization perspective i.e., focus on designing suitable optimization objectives or provide methods for efficiently solving such objectives. However, there are several critical applications where other issues such as obtaining good estimates of the uncertainty in an outcome, or data privacy constraints, are equally pressing. For example, in certain healthcare applications, patient data is private and uncertainty around prediction of health outcomes is required to manage risks. Similarly, for several applications in legal, finance, mission critical IoT systems (Victor et al., 2022) , etc., both the privacy and confidence guarantees are important for decision making.Since Bayesian learning methods are known for their ability to generalize under limited data, and can generally provide well-calibrated outputs (Wilson et al., 2016; Achituve et al., 2021a; b) , it makes them suitable for solving the challenges mentioned above. A straightforward Bayesian FL method would perform the following steps -each client does local posterior inference to obtain a distribution over weight parameters and then communicates the local posteriors to the server, the server receives the local posteriors from the clients and aggregates them to obtain a global posterior distribution which is then broadcast to the clients for the next round of training. However, this entire learning procedure is highly resource and communication-intensive. For a very simple example, solving an m-dimensional federated least squares estimation, this method will require O(m 3 ) computation on all the clients and server sites (Al-Shedivat et al., 2021) . This cost is much more as opposed to the standard FL which is generally O(m). How could the Bayesian methods be then used for FL settings without paying such high costs? Also, since clients can have varied degrees of computing resources, constraining each client to train identical models would be limiting. Then, how to aggregate the local posteriors when the clients are training personal models and the model weights across clients are not identical, becomes the next important question. We try to address these questions by proposing a framework that allows all clients to train their own personal Bayesian models (with varying model complexities), but still enables collaboration across clients by instilling information from the peer collaboration in the form of priors. This is achieved by using the functional space to indirectly determine priors on weights, as opposed to them being specified in the weight space as in traditional Bayesian methods.In addition, maintaining client data privacy is an important concern in FL. While attacks on neural networks that can recover training data from the models by either directly analyzing model parameters or indirectly by analyzing the model outputs are well known (Fredrikson et al., 2015; Papernot et al., 2017) , even in federated learning settings, when the data is restricted to the local client sites, it has been shown that FL methods are not immune to data leaks (Fowl et al., 2021; Wen et al., 2022; Fowl et al., 2023) . Therefore, to guarantee the privacy of the local client data, we apply a formal well-known standard of differential privacy (Dwork & Roth, 2014) . We use a carefully designed noise mechanism for data sharing that allows us to provide a privacy bound on the entire procedure.In this work, we propose a novel unified FL framework that can simultaneously address challenges arising due to limited data per client and heterogeneous compute resources across clients, together with the need for privacy guarantees and providing calibrated predictions. To the best of our knowledge, no previous work has jointly addressed all these learning issues in the FL context. Our positive results substantially increase the scope of FL for critical real-world applications.Our key contributions are summarized as follows :1. We propose a novel computation and communication efficient method for personalized federated learning based on Bayesian inference. In this method, collaboration across clients is achieved through a novel way of assigning prior distributions over the model parameters via the output space. We call this method Federated Bayesian Neural Networks (FedBNN).2. We present a non-trivial extension of our method that can work with heterogeneous settings where clients are not forced to train models of identical architectures. This provides more flexibility to the clients when the compute resources across clients are not identical or clients have existing pre-trained solutions.3. We provide a formal differential privacy guarantee for our method that applies to general settings irrespective of the client's learning algorithm and show that the method is able to learn effectively even under strict privacy guarantees.4. We evaluate our method on several datasets and show that it outperforms the baselines by a significant margin, particularly in data and model heterogeneous settings, making it particularly more useful for FL applications where a high degree of heterogeneity is naturally present. 
related work: 
 This section provides an overview of the most relevant prior work in the fields of federated learning, Bayesian FL, and Differential Privacy in FL.Federated Learning In the early 2000s, privacy preserving distributed data mining referred to training distributed machine learning models (Kargupta & Park, 2000; Gan et al., 2017; Aggarwal & Yu, 2008) , like distributed clustering (Merugu & Ghosh, Nov, 2003; 2005) , distributed PCA (Kargupta et al., 2001) , distributed SVMs (Yu et al., 2006) etc. Federated Learning was introduced as the FedAvg algorithm in the seminal work by (McMahan et al., 2017b) . Since then many different modifications have been proposed that tackle specific challenges. FedPD (Zhang et al., 2021) , FedSplit (Pathak & Wainwright, 2020), and FedDyn (Acar et al., 2021) proposed methods for finding better fixed-point solutions to the FL optimization problem. (Lin et al., 2020; Yurochkin et al., 2019; Wang et al., 2020; Singh & Jaggi, 2020; Chen & Chao, 2021) 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. FedDF (Lin et al., 2020) and (Li & Wang, 2019; Ozkara et al., 2021) achieve personalised FL solutions through collaboration by performing knowledge distillation on local client models. Personalised FL has been approached in many other ways like meta-learning (Fallah et al., 2020; Beaussart et al., 2021; Jiang et al., 2019; Khodak et al., 2019) , multi-task learning (Smith et al., 2017a; Li et al., 2021; Smith et al., 2017b) , by clustering the clients (Sattler et al., 2021b; Ghosh et al., 2020) and others (Collins et al., 2021; Li et al., 2020a; Yu et al., 2020; Shamsian et al., 2021) , (Wang et al., 2023; Makhija et al., 2022) , (Chen et al., 2022) uses Bayesian view based analysis to obtain a better trade-off between personal and global models. The personalised FL methods focus on improving the performance in the presence of statistical data heterogeneity across clients, but they do not work well when the size of dataset on the clients is limited.Bayesian Federated Learning The Bayesian approaches for federated learning can be broadly divided into two categories -methods using Bayesian inference for local learning at each client and methods that achieve collaboration via Bayesian mechanisms. Amongst the methods that use Bayesian approach for achieving collaboration across clients, FedBE (Chen & Chao, 2021) uses Bayesian mechanism to aggregate the locally trained neural networks to obtain a Bayesian ensemble at the server. (Bhatt et al., 2022) suggests using knowledge distillation and MCMC based method for training a global model. (Dai et al., 2020) on the other hand suggests the use of Bayesian Optimization and Thompson Sampling to obtain the solution to the global optimization problem. PFNM (Yurochkin et al., 2019) uses a Beta-Bernoullli process to obtain a global model from the local models but is only applicable to fully-connected networks, FedMA (Wang et al., 2020) extends PFNM for other types of networks. Recently, (Ozer et al., 2022 ) did an empirical study on various ways of aggregating variational Bayesian neural networks and their effects. On the other hand, FedPA (Al-Shedivat et al., 2021) was the first to use local Bayesian inference and suggested that the global posterior distribution of weights could be obtained by multiplying local posteriors and proposed an approximate and efficient way of computing local and global posteriors by using Laplace approximations with complexity that is linear in number of model parameters. This method focuses on obtaining a global solution and is less suited for the statistical heterogeneity present across clients (Cao et al., 2023) , and therefore we focus more on the methods that build personalized solutions for clients. Among such methods, pFedGP (Achituve et al., 2021b ) is a Gaussian Process-based estimation method where the kernel in these GPs is defined on the output of a neural network and is called Deep Kernel Learning. pFedGP works by collaboratively training a single deep kernel for all clients but using personalized GPs for prediction. The collaboration across clients while learning the global kernel is done as in FedAvg. FedLoc (Yin et al., 2020 ) also uses GP in FL but for regression tasks. pFedBayes (Zhang et al., 2022) uses variational inference for local posterior inference on Bayesian Neural Networks(BNNs) where the loss at each client is a combination of the data likelihood term and distance to the prior where prior distribution is replaced by the global distribution. The global distribution is obtained by aggregating the local prior distributions. FOLA (Liu et al., 2021) proposed using Laplace Approximation for posterior inference at both the server side and the client side. (Ozkara et al., 2023) and (Kotelevskii et al., 2022 ) also proposed methods that assume a hierarchical Bayesian model on the parameters but their main assumption is that the local parameters are obtained from the same distribution thus making them useful only in homogeneous settings. None of the methods described above explicitly handle heterogeneous settings. Moreover, for these methods choosing an appropriate prior distribution over the local model parameters is a challenge (Cao et al., 2023) . These issues led us to use functional space priors instead. Such priors have been studied in limited centralized settings (Tran et al., 2022; Sun et al., 2019; Flam-Shepherd, 2017) but not in FL settings.Differential Privacy in FL Since decentralized learning does not guarantee that the data will remain private, it is important that a 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 results (Dwork & Roth, 2014; Mironov, 2017; Kairouz et al., 2015) . For FL, (Geyer et al., 2017) and (McMahan et al., 2017a) independently proposed DP-FedSGD and DP-FedAvg algorithms which enhance FedAvg by adding Gaussian noise to the local client updates. Several other works focus on analyzing the privacy-utility trade-off in DP in FL setting (Ghazi et al., 2019; Girgis et al., 2021; Balle et al., 2019; Triastcyn & Faltings, 2019; Li et al., 2020b) . Recently, (Hu et al., 2020) proposed a DP-based solution for personalized FL which only works for linear models. And then (Noble et al., 2022) enhanced it for general models and heterogeneous data in FL. These methods, however, mostly focus on privacy guarantees while solving the FL optimization problem. They don't apply to general Bayesian settings or when substantial heterogeneity is encountered. 
proposed methodology: 
 We propose a novel method for Bayesian FL that enables privacy preserving personalised learning on clients under heterogeneous settings. The local training is done in a model agnostic way but collaboration amongst clients is enabled by passing information as the prior. And to achieve privacy, like most private algorithms, we use Gaussian mechanism to carefully add noise to the information sent outside a client. This method also produces well calibrated outputs. In this section, we first describe the problem statement and then go over the proposed solution.Consider an FL setting with N clients where each client i has local dataset X i of size n i drawn from the local data distribution D i . The goal of a personalised federated learning procedure is to obtain optimal weights for each client's local model, W i , using data X = N j=1 X j without actually accessing the data outside of each client, i.e., no other client can access any data in X i except the i th client itself but the clients could transfer knowledge via collaboration. In a personalised Bayesian learning procedure, the modified goal would be to learn distribution over local weights, IP(W i |X ) from X = N j=1 X j while still maintaining the client data privacy. However, the learning procedure faces challenges that are posed due to -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.e., n k >> n l for some clients k and l, and ii) the compute across clients is non-identical due to which it is not possible to train models of uniform architectures across clients, leading to non-identical weights, i.e., W i ̸ = W j for different clients i and j. Statistical heterogeneity implies that the data distribution across clients is non-IID.Here we propose an adaptive framework to learn personalised Bayesian Neural Network (BNN) based models for each client that can work effectively even under the challenges described above. The overall framework works iteratively in two steps -local optimization on the individual client to obtain local posterior distribution 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. Each of these two steps is described in detail below, and the detailed algorithm is given in the Appendix B.Local Setting Let each client in the network be training a personalised Bayesian NN, which for the client i is denoted by Φ i and is parameterised by weights W i . As commonly used in the literature, we assume that the individual weights of the BNN are Normally distributed and satisfy mean-field decomposition, i.e., w i,α ∼ N (µ i,α , σ 2 i,α ) for α ∈ [1, . . . , |W i |] where µ i,α is the mean of the gaussian distribution for the parameter α on the i th client and σ 2 i,α is the variance of the gaussian for the same parameter. To guarantee that σ i,α takes non-negative values for all clients i and all parameters α, we use a technique commonly used in inference procedures (Blundell et al., 2015) , wherein each σ i,α is replaced by another parameter ρ i,α during the training, such that σ i,α = log(1 + exp(ρ i,α )).We attain collaboration amongst clients under a mild assumption of availability of a general publicly accessible unlabelled dataset at the server. We call this dataset as Alignment Dataset (AD). This dataset is used for providing peer supervision to the clients by helping clients distill knowledge from other clients without sharing the model weight distributions. While in heterogeneous settings nonidentical architecture models across clients mean no direct way of aggregation and prior specification, even in homogeneous settings aggregating the weight distributions and specifying priors in the weight space is prone to errors due to reasons like non-alignment of weights across clients, insufficient understanding of the weight space, etc. Thus, for the purpose of collaboration, we move away from the weight-space to the function-space of the networks. 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 output Φ i (AD). The local output of the i th client is approximated by doing Monte Carlo sampling and drawing K weight samples, W(j) i : j ∈ [1, K], from its local posterior distribution IP(W i |X ).An aggregate of the obtained logits for these K weight samples under the client's own personal BNNmodel, Φ i (), is reported, i.e. Φ i (AD) = 1 K K j=1 Φ i (AD; W (j) i ).The obtained output for AD on each client is then sent back to server which forms an aggregated representation, denoted by Φ(AD), by doing a weighted aggregation of all clients' outputs, i.e., Φ(X) = N j=1 w j Φ j (X). The weights w's used in the aggregation could represent the strength of that particular client in terms of its data or compute resources, i.e., clients with high compute (or data) resources receive more weight as compared to 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 Alignment Dataset (AD) are included in the Appendix E.Prior Specification The Bayesian framework provides a natural way of incorporating supervision in the form of priors. Conventional methods in Bayesian deep learning provide priors over model parameters. However, the relationship between the model parameters 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 is 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 incorporate prior knowledge in 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 to achieve a desired functional output. While being more intuitive, these priors also help in instilling the prior external knowledge during the training of the neural networks.Local Optimization For the local optimization, the individual clients learn IP(W i |X i ) via variational inference. A variational learning algorithm, tries to find optimal parameters θ * of a parameterized distribution q(W i |θ) among a family of distributions denoted by Q such that the distance between q(W i |θ * ) and the true posterior IP(W i |X i ) is minimized. In our setting, we set the family of distributions, Q, to be containing distributions of the formw i,α ∼ N (µ i,α , σ 2 i,α ) for each parameter w i,α for α ∈ [1, . . . , |W i |].Let p(W i ; ψ) represent the prior function over the weights W i and is parameterized by ψ, then the optimization problem for local variational inference is given by :EQUATION= arg minEQUATIONFor inference in Bayesian neural networks, we use Bayes by Backprop (Blundell et al., 2015) method to solve the variational inference optimization problem formulated above.At the beginning of each local optimization procedure (in each global communication round a specific client is selected), we use the global information obtained from the server Φ(AD) to intialize the prior for the BNN. Specifically, at the beginning of each local training round, the selected clients first tune their priors to minimize the distance between the local output, Φ i (AD; W i ) 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 strictly precise for the local model optimization, we consider this aggregated output as "noisy" and correct it before using for optimization. Specifically, we generate Φ corrected i as a convex combination of the global output and the local output for a tunable parameter γ. For the i th client,Φ corrected i = γ Φ(AD) + (1 -γ)Φ i (AD; W i ).(3)The prior optimization steps then optimize the distance between Φ corrected i and Φ i (AD; W i ) to train the prior parameters ψ, with the aim of transferring the global knowledge encoded in Φ corrected i to the local model. Precisely,EQUATIONWhen the outputs Φ(X; W) are logits, we use cross-entropy or the negative log-likelihood loss as the distance measure. The optimization involves training the client's personal BNN Φ i to only learn the parameters of the prior distribution denoted by ψ. This way of initializing the BNN prior enables translating the functional properties, as captured by Φ i (AD; W i ), 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 to find the best q(W i |θ) fixing the prior distribution through the following optimization problem :θ * i = arg min θ:q(Wi|θ)∈Q KL[q(W i |θ)||p(W i ; ψ * i )] -E q(Wi|θ) [logIP(X i |W i )].(5)In this method, we measure the privacy loss at each client. To control the release of information from the clients, we add a carefully designed Gaussian mechanism wherein we add Gaussian noise to the Φ i (AD) that is being shared by each client. Specifically, each client i uploads Φ i (AD) DP = Φ i (AD) + N (0, σ 2 g ) to the server and then the server aggregates Φ i (AD) DP across clients to obtain and broadcast Φ(AD) DP which is used by the clients in their next round of local optimization. The variance of the noise depends on the required privacy guarantee. 
privacy analysis: 
 Since our algorithm does not specifically share model weights outside of the clients but shares the client models' output on a public dataset, it might seem intuitive that the algorithm is private but knowing that privacy is lost in unexpected ways, we present a formal Differential Privacy based guarantee on the algorithm. Our analysis in this section focuses on providing record-level DP guarantee over the entire dataset X . This analysis quantifies the level of privacy achieved towards any third party and an honest-but-curious server. In this section we first define DP and then present our analysis. Additional definitions and the results used in the proof are mentioned in the Appendix C. Definition 4.1 ((ϵ, δ)-Differential Privacy). A randomized algorithm M is (ϵ, δ)-DP if for any two neighboring datasets D and D ′ that differ in at most one data point, the output of the algorithm M on D and D ′ is bounded asIP[M(D) ∈ S] ≤ e ϵ IP[M(D ′ ) ∈ S] + δ, ∀S ⊆ Range(M).Due to the lack of space, definitions of zero-concentrated differential privacy (zCDP), L 2 Sensitivity of an algorithm, and results corresponding to Gaussian Mechanism, Sequential Composition and Parallel Composition of privacy are given in Appendix C. Building on top of these results, we now state the privacy budget of our algorithm. Theorem 4.2 (Privacy Budget). The proposed algorithm is (ϵ, δ)-differentially private, if the total privacy budget per global communication round per query is set toρ = ϵ 2 4EKlog 1δ 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 Φ i (AD), the local mechanism at each client becomes ρ-zCDP for ρ = ∆ 2 2σ 2 . Since each client outputs the logit representation for each input, i.e., the normalized output of the clients, ∆ 2 ≤ 2. Suppose in each global communication round we make K queries to each client, then by sequential composition C.2, we get EKρ, for E number of global communication rounds. By parallel composition C.3, the total privacy loss for all N clients is the maximum of the loss on each client and therefore remains EKρ. Relating it to the (ϵ, δ)-DP from C.1, we get ρ = ϵ 2 4EKlog 1 δ for any δ > 0.Our analysis does not assume any specifics of how each client is trained and is therefore applicable in more general settings. However, we present a pessimistic analysis by providing a worst-case analytical bound. This is because we assume that a change in single data point may entirely change the output of the algorithm and upper bound the L 2 sensitivity ∆ 2 ≤ 2, and also, since the public dataset remains common throughout the rounds, the actual privacy loss due to querying on the public dataset will 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 . 
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, we experiment with three types of heterogeneity -i) heterogeneity in data resources (amount of data), ii) heterogeneity in compute resources, and iii) statistical heterogeneity (non-IID 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 showing calibrated outputs are included in the Appendix E, G and D respectively.Datasets We choose three different datasets commonly used in prior federated learning works from the popular FL benchmark, LEAF (Caldas et al., 2019) including MNIST, CIFAR-10 and CIFAR-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-100 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.Simulation Details We simulate three different types of heterogeneous settings -corresponding to heterogeneity in 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 a 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 to the other high-capacity 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 a ReLU and 2 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 baselines for a fair comparison. We include the results of our 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 2500 data instances per class respectively for training. We simulate statistical heterogeneity by creating non-IID data partitions across clients. We work in a rather strict non-IID 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 5 of the 10 classes for MNIST and CIFAR-10, and 20 out of 100 classes for CIFAR-100.We run all the algorithms for 200 global communication rounds and report the accuracy on the test dataset at the end of the 200 th 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 a fixed number of epochs, which is kept to 50 in experiments, before entering the collaboration phase. The hyper-parameters 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.0001 and run the prior optimization procedure for 100 steps. Then with the optimized prior we train the local BNN using Bayes-by-Backprop, with Adam optimizer, learning rate = 0.001 and batch size = 128. The noise effect γ is selected after fine-tuning and kept to be 0.7. For these experiments, the aggregation weight w j for each client j used to compute Φ(X) is set to 1/N . All the models are trained on a 4 GPU machine with GeForce RTX 3090 GPUs and 24GB per GPU memory. For evaluation, we report the classification accuracy obtained by running the trained models on test datasets from the MNIST, CIFAR10 and CIFAR100 datasets.Baselines We compare our method against the standard 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 (per round ϵ < 0.1) to provide perspective on the privacy. The FL baselines include -i) FedAvg, the de-facto FL learning algorithm that trains a global model, ii) FedProx, an enhancement of the FedAvg algorithm in the presence of statistical heterogeneity across clients giving a global model, iii) pFedME, which uses personalized models on each client using Monreau envelopes in loss. The Bayesian baselines include -i) pFedGP, a 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) DP-FedAvg, the FedAvg algorithm with gradient clipping and noise addition to the gradient at each client. For all the experiments, the hyper-parameters were obtained by tuning on a held-out validation dataset. We used our own implementation of the pFedBayes algorithm since the source code was not publicly available.The performance of our method and the baselines under the non-IID data setting are reported in Table 1 .Under the non-IID setting, we report the results corresponding to different dataset sizes on each client. To recall, in the small, medium, and full settings, each client has access to 50, 100, and 2500 training data points per class respectively. We observe that our method with homogeneous 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 1 compares the performance of our method with the highest-performing 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 to gain about 10% increase in their performance. Also, the performance degradation of our method with a change in the number of clients with limited data resources is 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 (Sattler et al., 2021a) .The superior performance of our method could be attributed to the effective collaboration achieved via the proposed framework, wherein in each communication round, instead of aggregating clients' weight distributions and using these aggregated weights for initialization at each client, we achieve collaboration by first distilling peer knowledge in the form of the aggregated output on the AD, Φ(AD), and then ensuring that this knowledge is successfully transferred to each client by specifying priors in the functional-space of the client model. Furthermore, the parameter in Equation 3 allows the clients the flexibility to choose the amount of global knowledge that needs to be incorporated providing flexibility on the degree of personalization. 
discussion: 
 We propose a novel method for personalized Bayesian learning in heterogeneous federated learning settings and demonstrate that our method is able to outperform the baselines under different types of heterogeneous situations, while also providing a privacy guarantee and calibrated responses. The experiments show that the method is particularly useful for clients with lower data and lower compute resources as they can 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, it is typically a very mild requirement as such data can often be obtained from several open sources on the web. And with the recent advances in the generative AI and it's use for generating training data, it would be interesting to see if these methods can be used to readily provide application-specific alignment datasets in the future work. The privacy analysis on the method provides an intuitive and a rigorous guarantee with various tunable knobs that can be adjusted to achieve the desired privacy-utility trade-off. While the application explored in the proposed work consists of image related tasks, both the proposed framework and the privacy analysis is generic and independent of specific training algorithms, i.e., it provides a way for personalized FL for general tasks, therefore resulting in its wide applicability in various applications across data modalities. The recent use of transformer based Bayesian methods Zheng et al. (2023) ; Maged & Xie (2022) ; Tian et al. (2023) in varied applications also provides more evidence on the applicability of the proposed framework for settings where much larger neural networks are 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.Supplement for "A Bayesian Approach for Personalized Federated Learning in Heterogeneous Settings"In this supplementary material, we 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 in the privacy analysis of our method. 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 and the communication and computation cost of the procedure.Consider a learning setting where we are trying to train a neural network on a dataset X . The aim of this setting is thus to obtain the set of weights, denoted by W, for the corresponding neural network that best fits the data. We could also view a neural network as a model that outputs IP(y|x, W) which is the distribution of the label y for a given data point x 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 a given set of datapoints X = (x i , y i ) n i=1 by solving the following optimization problem.W M LE = arg max W i logP (y i |x i , W)This optimization could be solved by gradient descent based methods and obtains a point estimate of the weight vector, denoted by W M LE .The Bayesian learning methods, on the other hand, obtain a posterior distribution on the weights given the training data, IP(W|X ) as opposed to the point estimates. The predictions for any new data point, x, are then obtained by taking expectation of the prediction under the posterior distribution, y = E w∼IP(W|X ) [IP(y|x, w)]. Exact inference of the posterior distribution, however, is intractable for neural networks. Variational inference is 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 (Hinton & van Camp, 1993) . Specifically, variational inference tries to learn a simpler parameterized distribution q(w|θ) from a family of distributions Q by optimizing the parameters θ such that the new distribution q(w|θ * ) obtained for the optimal value of θ is close to the true posterior distribution IP(W|X ). Precisely, the optimization problem looks likeEQUATION= arg minEQUATION= arg minEQUATIONwhere p(W; ψ) signifies the prior distribution over weights W parameterized by ψ. The prior distribution is typically used to encode any previously available information about the weights of the network. The above given objective is the same objective as in Equation 1 that is used for local training in our method.The pseudo-code of the algorithm used in the FedBNN method is included in the Algorithm 1. The Algorithm 1 works in the setting when there is a server connected to N clients with each client i having local dataset X i of size n i drawn from the local data distribution D i , and the server has an auxilliary unlabelled dataset called AD. The output of the algorithm is the set of personalized models Φ i parameterized by W i for each client i. All W i 's, instead of being point estimates, are determined by a posterior distribution IP(W i |.) which is learnt from the data via variational inference.As mentioned in the Section 3.2, the learning procedure first optimizes the prior parameters by minimizing Equation 4 and then learns the posterior parameters keeping the prior fixed by minimizing Equation 5.Algorithm 1 FedBNN Algorithm Input: number of clients N , number of global communication rounds T , number of local epochs E, weight vector [w 1 , w 2 , . . . w N ], noise parameter γ Output:Personalised BNNs {Φ i |i ∈ [1, N ]}, parameterized by W i ∼ IP(W i |X ) Server Side - X = AD for t = 1 to T do Select a subset of clients N t for each selected client i ∈ N t do Φ i (X) = LocalTraining(t, Φ(X) (t-1) , X) end for Φ(X) (t) = N j=1 w j Φ j (X) end for Return Φ 1 (T ), Φ 2 (T ) . . . Φ N (T ) LocalTraining(t, Φ(X) (t-1) , X) Run inference on X to obtain Φ i (X) Generate Φ corrected i (X) = γ Φ(X) (t-1) + (1 -γ)Φ i (X) for each prior epoch do Minimize CrossEntropy(Φ corrected i (X), Φ(X) (t-1) ) to obtain prior parameters ψ of the BNN Φ i end for for each local epoch do Minimize KL[q(W i |θ)||p(W i ; ψ * )] -E q(Wi|θ) [logIP(X i |W i )] over {θ : q(W i |θ) ∈ Q} to obtain θ * end for IP(W i |X ) ≈ q(W i |θ * ) Obtain K Monte-carlo samples W (j) i : j ∈ [1, K] from IP(W i |X ) Compute Φ i (X) = 1 K K j=1 Φ i (X; W (j) i ) Return Φ i (X)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.A generalization of differential privacy is concentrated differential privacy(CDP). And an alternative form of concentrated differential privacy called zero-concentrated differential privacy(zCDP) was proposed to enable tighter privacy analysis (Bun & Steinke, 2016) . We will also use the zCDP notion of privacy for our analysis. The relationship between standard DP and zCDP is shown below. Proposition C.1 ((ϵ, δ)-DP and ρ-zCDP). For a randomized algorithm M to satisfy (ϵ, δ)-DP, it is sufficient for it to satisfy ϵ 2 4log 1 δ -zCDP. And a randomized algorithm M that satisfies ρ-zCDP, also satisfies (ϵ ′ , δ)-DP where ϵ ′ = ρ + 4ρlog 1 δ .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 (Yu et al., 2019) and is recalled below. Proposition C.2 (Sequential Composition). Consider two randomized mechanisms, M 1 and M 2 , if M 1 is ρ 1 -zCDP and M 2 is ρ 2 -zCDP, then their sequesntial composition given by (M 1 (), M 2 ()) is (ρ 1 + ρ 2 )-zCDP.Proposition C.3 (Parallel Composition). Let a mechanism M consists of a sequence of k adaptive mechanisms, (M 1 , M 2 , . . . M k ) working on a randomized partition of the D = (D 1 , D 2 , . . . D k ), such that each mechanism M i is ρ i -zCDP and M t :t-1 j=1 O j × D t → O t , then M(D) = (M 1 (D 1 ), M 2 (D 1 ), . . . M k (D k )) is max i ρ i -zCDP.After computing the total privacy loss by an algorithm using the tools described above, we can determine the variance of the noise parameter σ for a set privacy budget. The relationship of the noise variance to privacy has been shown in prior works by (Dwork & Roth, 2014; Yu et al., 2019) and is given below. Definition C.4 (L 2 Sensitivity). For any two neighboring datasets, D and D ′ that differ in at most one data point, L 2 sensitivity of a mechanism M is given by maximum change in the L 2 norm of the output of M on these two neighboring datasets∆ 2 (M) = sup D,D ′ ||M(D) -M(D ′ )|| 2 .Proposition C.5 (Gaussian Mechanism). Consider a mechanism M with L 2 sensitivity ∆, if on a query q, the output of M is given as M(x) = q(x) + N (0, σ 2 ), then M is ∆ 2 2σ 2 -zCDP.Model calibration is a way to determine how well the model's predicted probability estimates the model's true likelihood for that prediction. Well-calibrated models are much more important when the model decision is used in critical applications like health, legal etc. because in those cases managing risks and taking calculated actions require a confidence guarantee as well. Visual tools such as reliability diagrams are often used to determine if a model is calibrated or not. In a reliability diagram, model's accuracy on the samples is plotted against the confidence. A 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 whereas MCE measures the maximum deviation between the accuracy and the confidence. The definitions and empirical formulas used for calculating ECE and MCE are as given below.ECE = E P [IP( Ŷ = Y | P = p) -p], MCE = max p∈[0,1] |IP( Ŷ = Y | P = p) -p|. Empirically, ECE = M i=1 |B i | n |accuracy(B i ) -confidence(B i )|, MCE = max i∈[1,M ] |accuracy(B i ) -confidence(B i )|,where B i is a bin with set of indices whose prediction confidence according to the model falls into the range i-1 M , i M . Figure 2 shows the reliability diagram along with the ECE and MCE scores for our method measured on MNIST and CIFAR-10 dataset in the non-IID data setting.In FedBNN, the alignment dataset (AD) is used to achieve collaboration across clients. Since the only assumption on AD is for it to be of the same domain as the target application, there is no practical constraint on obtaining the AD in real-world settings. In many cases it could be obtained from web, for example images from common datasets in Huggingface, texts from Wikipedia, Reddit etc. The use of AD is not different from how several other methods use an additional dataset for augmentation. The effect of size of AD on the performance of models is demonstrated in Figure 3 for CIFAR-10 dataset in the small data and non-IID setting. In that figure, we observe that when the size of AD is 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 2 . We run these experiments on 20 clients for CIFAR-10 dataset where each client had access to only 5 of the 10 classes and each client belonged to the medium data setting. For the first experiment, we use a held-out dataset from the CIFAR-10 data as AD but vary the composition of the dataset by changing the distribution of the classes present in the AD, for example, CIFAR10( 10) is composed of all 10 classes present in the CIFAR-10 dataset but CIFAR10(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 a significantly different dataset SVHN consisting of the colored house number images is used. Table 2 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 capacities 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 parameter γ controls the amount of global knowledge to be incorporated on each client and helps in regularization of the client models which enables them to generalize better and perform better on the test data. 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 a publicly available dataset, AD could be transmitted to the clients by specifying the source and the indices, and does not really needs to be 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 3 , 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 that transmit model parameters in each communication round, making it practically more useful.Computation Cost Similarly, the computation cost of 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 on the AD (with size 2000), thus the server side computation cost in our method is many orders of magnitude lower than the conventional methods and does not 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 a Bayesian Neural Network at the client side, which in turn depends on the 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 to obtain the posterior distributions where the posterior distributions are characterized by the mean and standard deviation. A re-parameterization trick is used to compute the mean and std of the distributions from the backpropagated gradients. Thus the cost for obtaining the posterior distributions is similar to the cost of backpropagation. Moreover, since the method only uses gradient updates, the optimizations used for SGD like asynchronous SGD etc. could be readily used for obtaining the posteriors. An unrelated but similar algorithm in (Hernández-Lobato & Adams, 2015) does probabilistic backpropagation to 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 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 the inference procedure used to obtain the posterior distributions and for which Bayes by Backprop provides an efficient mechanism. Several works in the recent past have discussed the use of related Bayesian inference based methods for training uncertainty-aware transformers (Zheng et al., 2023; Tian et al., 2023; Maged & Xie, 2022) proving that Bayesian methods are not limited to use in simpler models. And therefore, our framework can also be extended to apply in settings where much larger neural networks are required. Privacy vs Performance Since the amount of noise required to be added to the client's outputs via the Gaussian Mechanism is directly proportional to the guaranteed privacy, we test the affect of the privacy guarantee on the performance of the proposed framework by comparing the performance of the method with varying ϵ and δ = 10 -4 . The results are reported in Table 3 . We observe that, as expected, when we reduce the amount of 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 not drastic as the clients can tune the level of personalization or global knowledge required by appropriately setting the parameter γ in Equation 3.To test the performance of the proposed method when a large number of clients are involved in the setup, we did additional experiments with 500 clients and non-IID setting with 5 classes per client in the medium data setting on the CIFAR-10 dataset where in each communication round only 10% of the clients are selected for participation and γ = 0.7. The obtained results at the end of 200 th communication round are reported in Table 4 . We observe that the homogeneous version of our method is better than the baselines by a significant margin and the heterogeneous version is slightly better than the baselines. 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_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. 

The 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.

----
Post Rebuttal

The rebuttal addresses my concerns on the DP guarantees, as such I raise my score to 6.
strength_weakness: strength_weakness: * 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. 
* 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*).

**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.**
questions: 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: limitations: FedBNN relies on the availability of a public auxiliary dataset.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. 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.
questions: 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: limitations: N/A
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
[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.
questions: questions: See Weaknesses above.
limitations: limitations: NA
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. 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.
3. 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.
4. 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.
questions: questions: 1. Please introduce the difficulties and improvements in integrating the Bayesian network into the data distillation-based federated learning framework.
2. What is the performance of this method under other data heterogeneity or model structures?
limitations: limitations: The authors discussed possible social impact and application and also touched on the potential limitations of the method.

Comments: title: Thank You
Response by reviewer
Kindly Request for Reviewer's Feedback
Kindly Request for Reviewer's Feedback
Kindly Request for Reviewer's Feedback
Kindly Request for Reviewer's Feedback
Response to Official Comment by Reviewer PtiT
comment: Dear Reviewer yjRB,

Thanks for maintaining positive score of 6 of our paper. 

Please let us know if you still have any other questions.

Best,

The Authors
Thank the authors for providing further clarification. 

I have read the rebuttal and other reviews, I decide to maintain my original score of 6.
Dear Reviewer PtiT,

We sincerely appreciate the time you have taken to provide feedback on our work, which has helped us greatly improve its clarity, among other attributes. This is a gentle reminder that the Author-Reviewer Discussion period ends at 11:59 pm AoE on August 13. We are happy to answer any further questions you may have before then, but we will be unable to respond after that time.

If you agree that our responses to your reviews have addressed the questions you listed, we kindly ask that you consider whether raising your score would more accurately reflect your updated evaluation of our paper. Thank you again for your time and thoughtful comments!

Sincerely,

Authors
Dear Reviewer 4fxY,

We sincerely appreciate the time you have taken to provide feedback on our work, which has helped us greatly improve its clarity, among other attributes. This is a gentle reminder that the Author-Reviewer Discussion period ends at 11:59 pm AoE on August 13. We are happy to answer any further questions you may have before then, but we will be unable to respond after that time.

If you agree that our responses to your reviews have addressed the questions you listed, we kindly ask that you consider whether raising your score would more accurately reflect your updated evaluation of our paper. Thank you again for your time and thoughtful comments!

Sincerely,

Authors
Dear Reviewer yjRB,

We sincerely appreciate the time you have taken to provide feedback on our work, which has helped us greatly improve its clarity, among other attributes. This is a gentle reminder that the Author-Reviewer Discussion period ends at 11:59 pm AoE on August 13. We are happy to answer any further questions you may have before then, but we will be unable to respond after that time.

If you agree that our responses to your reviews have addressed the questions you listed, we kindly ask that you consider whether raising your score would more accurately reflect your updated evaluation of our paper. Thank you again for your time and thoughtful comments!

Sincerely,

Authors
Dear Reviewer rtXt,

We sincerely appreciate the time you have taken to provide feedback on our work, which has helped us greatly improve its clarity, among other attributes. This is a gentle reminder that the Author-Reviewer Discussion period ends in just around 12 hours from this comment, i.e., 11:59 pm AoE on August 13. We are happy to answer any further questions you may have before then, but we will be unable to respond after that time.

If you agree that our responses to your reviews have addressed the questions you listed, we kindly ask that you consider whether raising your score would more accurately reflect your updated evaluation of our paper. Thank you again for your time and thoughtful comments!

Sincerely,

Authors
Dear Reviewer PtiT,

We thank you for your detailed response to our rebuttal and sincerely appreciate your thoughtful comments. We further hope to effectively address your concerns  on novelty and additional experiments through the points we make below.

-- *Additional Experiments*

We are glad to hear that you believe the additional experimental results enhance the paper and we thank you again for suggesting these new experiments. Although the limited time during the rebuttal period prevented us from completing the additionally suggested experiments on more datasets and settings, we want to assure you that we have already established the necessary experimental setup which will allow us to easily incorporate results on more datasets and clients in the next revision, and we are fully committed to doing so. We appreciate your understanding and look forward to sharing these updated findings in the final submission.



-- *Novelty*

We apologize if it seemed we did not fully address your concerns about the novelty of our work. We would like to take this opportunity to provide further clarification.

The key motivation of this work is to enhance the applicability of FL across diverse real-world settings by addressing several interconnected and prevalent challenges in federated environments. These challenges include - i) the issue of small data volume per client, ii) compute heterogeneity across clients, (both of which are known common limitations in federated learning environments), iii) the need for uncertainty quantification to measure and manage uncertainty in predictions, which is crucial for critical applications like in healthcare and legal domains, and iv) the need for privacy preservation which is essential for ensuring the confidentiality of sensitive information and is also important in critical domains. To tackle these issues, we have developed an integrated end-to-end framework, in which each client performs local training using Bayesian learning with its own tailored BNN architecture (and that could be different across clients), and the collaboration is achieved by the novel mechanism of knowledge distillation through prior specification. We also emphasize on the fact that the conventional Bayesian learning methods, whether centralized or federated, often focus on priors in the weight space, which makes them restricted for utilization in FL settings where the clients have heterogeneous data and compute resources, and is also communicationally intensive. Thus, our framework is designed to effectively address these issues as well—issues that would arise from merely applying conventional Bayesian learning methods in FL—through a novel collaboration mechanism. 

We recognize that some components of our approach may seem individually straightforward when viewed in isolation. However, it’s crucial to understand that our approach is not just a collection of techniques but a carefully designed system that addresses multiple challenges in FL. It's appropriate formulation into a cohesive and effective framework, without introducing additional complications, for personalized federated learning represents a significant advancement in the field. We hope this further explanation clarifies the novelty and impact of our contributions, and we appreciate your thoughtful consideration of our response. 

In essence, the key contributions of our work are not just in - *"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."*. We thank the reviewer for highlighting the importance of this discussion in the paper, and we will use the additional space available in the final manuscript to include this discussion.
Overall, I believe the authors' rebuttal has provided some valuable new information. My updated comments are as follows:

**Experiments**  
I appreciate the additional experimental data, which enhances the paper. However, the scope remains somewhat limited, with one dataset and a small number of clients. While I acknowledge the authors' efforts during the rebuttal period, I believe the experiments could be further strengthened.

**Model Structure and Presentation Issues**  
Thank you for the clarifications. I apologize for any earlier misunderstandings. These explanations are helpful, and I look forward to seeing them incorporated into future revisions of the paper.

**Novelty**  
I feel that the authors did not directly address my concerns regarding the novelty (e.g., concern 1). While I agree that some of the innovative aspects highlighted in the rebuttal are indeed non-trivial and novel to some extent, as I mentioned in my first weakness point, there are two issues: 1) Each point, when considered individually, lacks sufficient insight or difficulty, and the approach seems relatively straightforward (in my opinion); 2) When taken together, these innovative points do not seem to be well connected, making the overall contribution less compelling.

In summary, while some concerns have been addressed, others remain unresolved, leading me to adjust my score from 3 to 4. I hope this feedback is helpful as the authors continue through the review process.