A Bayesian Approach for Personalized Federated Learning in Heterogeneous Settings

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.

We address reviewer's specific concerns below -

-- Clarification on the contribution

We 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 -

1. 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.
2. 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.
3. 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.

-- On model structures in Experiments

We 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".

-- On different data heterogeneity in Experiments

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

-- Presentation Issues

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

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

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.

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

We address reviewer's key concerns below -

-- Differentitaion with FedPop

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

1. 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.
2. 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.
3. 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.

-- Privacy Analysis

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

-- Functional space priors versus traditional weight space priors

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

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

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

We address reviewer's key concerns below -

-- "Bayesian frameworks and knowledge distillation have been independently integrated into federated learning and studied by many works. "

We 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 -

1. 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.
2. 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.
3. 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.

-- "Experimental Results"

We 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).

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

Dear Reviewer rtXt,

We are thankful to the reviewer for the thorough feedback on our work.

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

Below we address the reviewer's concerns on the DP analysis -

-- Dependence on $\sigma$

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

-- Clarification on the bound on $\Delta$

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

$\Delta (\Phi) = || \Phi(D) - \Phi(D')||_2$

Since all other data-points between $D$ and $D'$ are identical, the $L_2$ sensitivity of $\Phi$ becomes -

$\Delta (\Phi) = || \Phi(D(i,:)) - \Phi(D'(i,:))||_2$

Now, $\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$).

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

Best,

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