Personalized Federated Learning of Probabilistic Models: A PAC-Bayesian Approach

3. The novelty of the proposed method PAC-PFL is limited since it seems that the PAC-PFL only combines some techniques, such as PAC-Bayesian, FedAvg, and SVGD.

While it is true that our method hinges on the existing pillars such as the PAC-Bayesian theory, FedAvg, and SVGD, there are substantial differences between our approach and each of these methods that we would like to highlight in the following.

1. Innovative formulation: We propose a rigorous framework for personalized learning by formulating a shared hyper-posterior and regarding posterior inference by each client as the personalization step (refer to Fig. 1 in the paper). Unlike other federated learning papers that employ PAC bounds solely for analyzing generalization, our approach uniquely integrates the PAC-Bayesian theory into algorithm design. Particularly, we specify the hyper-posterior by minimizing a PAC-Bayesian bound. We believe this formulation proposes an original and innovative framework.
2. Theoretical advances: Under this formulation, we establish a novel PAC-Bayesian bound at the server level in Theorem 4.1. This bound isn't a straightforward deduction from existing bounds, as reflected in our proof. The PAC-Bayesian theory has been classically developed for a single dataset of i.i.d. samples. However, our framework deals with multiple datasets from heterogeneous clients, necessitating the derivation of non-trivial extensions of the existing PAC-Bayesian bounds. These extensions, presented in our paper, require entirely new proofs and are absent in current literature.
3. Implementation: In our practical algorithm, we employ SVGD as a method to approximate the hyper-posterior. However, the critical aspect lies in determining the relevant probability distribution to approximate, which is the optimal hyper-posterior specified by our formulation and theoretical analysis. The structure of our algorithm for learning the shared hyper-posterior resembles FedAvg, where clients send the gradients of a cost function with respect to their local data and the server aggregates these gradients for updating the shared model. However, the crucial distinction lies in defining the cost to minimize and accordingly, determining the computation of gradients which is enabled through our formulation and theoretical analysis. Finally, FedAvg lacks the personalization aspect that is fundamental in PFL and is realized through posterior inference in our framework.

In conclusion, the combination of these tools for achieving a PFL method, in our opinion, is far from being obvious.

4. The related work [3] for tackling FL for small datasets is omitted.

The authors of [3] profoundly explore FL in scenarios where clients' datasets are small. However, they concentrate on training a global shared model within their framework, without incorporating personalization. We highlight that personalized FL (PFL) is the main focus of our paper. Given the substantial body of work in PFL, our literature review has predominantly delved into personalized approaches.

We, however, recognize the relevance of [3] to our work, particularly in considering clients with small datasets. In the final version of the paper, we will cite this paper in the introduction. Thank you for bringing this reference to our attention.

5. For experiment results, almost all FL methods usually outperform the Pooled GP baseline, which is strange and should be further explained.

Please refer to our answer 3.1 to reviewer D4J2.

6. In experiment results of tables and figures, the reported baselines are generally different from each other, which is confusing.

Thank you for this comment. We fully agree that baselines across different experiments should be consistent as much as possible. However, they cannot be exactly the same, as detailed in the sequel. There are three primary differences between the baselines used in regression and classification. Firstly, pFedBayes is exclusively applied to the FEMNIST dataset, as explained in Section 8.6.1, because the authors of pFedBayes only consider classification tasks in their experiments. Second, we did not apply FedAvg to the PV dataset due to its anticipated suboptimal performance stemming from the dataset's high level of heterogeneity. Finally, following your comment and the comment from reviewer D4J2, we realized that Pooled BNN was missing on the FEMNIST dataset and added it to Table 6 of the paper. Particularly, the accuracy of the Pooled BNN model is $89.4$% when each client has $20$ training samples and $94.3$% when each client uses all the available data. The respective calibration errors are $0.106$ and $0.073$. We believe that this will make the baselines more homogeneous. In summary, with reference to the new version of the paper, pFedGP, MTL, MAML, Vanilla, and Pooled baselines are consistent across all experiments, while pFedBayes and FedAvg are exclusively employed for classification tasks.

Dear reviewer 5CVc,

We appreciate your valuable feedback and thoughtful comments. Below, please find our answers to all your questions and concerns. We have submitted a revised version of the paper, encompassing the majority of the reviewers' comments. The main changes in the text appear in blue fonts. Should the paper be accepted, a comprehensive updated version will subsequently be submitted, addressing any outstanding comments. These updates, while relatively straightforward, were not included during the rebuttal phase due to timing constraints and are highlighted within specific responses.

1. The description of uncertainty calibration is insufficient which is unfriendly for new readers. The related works[1-2] for tackling uncertainty calibration are omitted.

We use calibration error as a metric for evaluating our trained model. Since the computation of this metric is not a central focus of our paper, we have restricted our discussion to citing relevant works (Rothfuss et al., 2021; Kuleshov et al., 2018), from which we derived the necessary formulas. The calibration formula employed for classification is consistent between (Kuleshov et al., 2018) and the first reference you provided, (Guo et al., 2017). However, it is worth noting that (Kuleshov et al., 2018) additionally present a metric for calibration error in regression, which prompted us to cite this work.

We acknowledge that including a subsection in the appendix containing the formulas for computing the calibration error will enhance the completeness and self-sufficiency of our exposition. We will incorporate this in the final version and cite the references that you provided. We appreciate your feedback, and we hope that this enhancement will make our work more clear.

2. The studied issues, such as small datasets, highly heterogeneous data distribution, uncertainty calibration, and new clients, should be further organized and summarized.

We thank you for your comment and recognize that the addressed challenges might appear scattered over the introduction. We address four main challenges, which are: (c1) quantifying epistemic uncertainty, (c2) handling highly heterogeneous clients, (c3) considering clients possessing small datasets, and (c4) accounting for collecting new data over time. To enhance clarity, we are now enumerating these challenges in a paragraph in the introduction, assigning them labels (c1)-(c4). These labels are consistently used throughout the paper when referring to these challenges. Furthermore, in the experiments section, we now explicitly demonstrate how each experiment relates to these challenges. We invite you to refer to the updated version of the paper for further details.

Dear reviewer HvmF,

We appreciate your valuable feedback and thoughtful comments. Below, please find our answers to all your questions and concerns. We have submitted a revised version of the paper, encompassing the majority of the reviewers' comments. The main changes in the text appear in blue fonts. Should the paper be accepted, a comprehensive updated version will subsequently be submitted, addressing any outstanding comments. These updates, while relatively straightforward, were not included during the rebuttal phase due to timing constraints and are highlighted within specific responses.

1. How is the hyper-prior formulated? We select a multi-variate Gaussian distribution with a diagonal covariance matrix as the hyper-prior, as indicated in Table 1 in Appendix 8.3.6. In the main paper, we refer to this table, mentioning that it contains selection guidelines in Section 5 after presenting the federated algorithm.

To further clarify the hyper-prior selection in the main part of the paper, we have included this information regarding the hyper-prior selection in Section 5, specifically in the paragraph discussing the federated algorithm (please see the updated version of the paper). We appreciate your feedback.

2. Is it possible to claim differential privacy for a finite number of SVGD updates?

As detailed in Section 5 within the paragraph labeled "SVGD at the server level," since SVGD is a deterministic approach, it compromises the inherent privacy of the optimal hyper-posterior established in Lemma 4.3. However, we address this by suggesting a privacy-preserving variant of PAC-PFL in Appendix 8.4 for a finite number of SVGD updates. To enhance clarity and better emphasize our proposal of the privacy-preserving algorithm in the appendix, we have revised the paragraph titled "SVGD at the server level". Please refer to the updated paper for further details. Thank you for bringing this to our attention and we hope the updates provide greater clarity.

Dear reviewer QnoR,

We appreciate your valuable feedback and thoughtful comments. Below, please find our answers to all your questions and concerns. We have submitted a revised version of the paper, encompassing the majority of the reviewers' comments. The main changes in the text appear in blue fonts. Should the paper be accepted, a comprehensive updated version will subsequently be submitted, addressing any outstanding comments. These updates, while relatively straightforward, were not included during the rebuttal phase due to timing constraints and are highlighted within specific responses.

1. As for baselines, only 1 and the latest of them are proposed in 2022, methods that were proposed in 2023 should also be considered.

In the submitted manuscript, we focused on comparing our method with personalized approaches for learning probabilistic models. It is conceivable that we might have overlooked some methods that were published while we were finalizing the paper. In conducting a literature review following your comment, we encountered (Ozkara et. al., 2023), which introduces a PFL approach with an interpretation within the empirical Bayes framework. However, to the best of our understanding, this algorithm does not offer uncertainty quantification and lacks a specific discussion on clients with small datasets. Upon acceptance, we intend to cite this paper in the introduction and highlight the differences from our work. If you are aware of any other methods, we would greatly appreciate it if you could inform us.

(Ozkara et. al., 2023) Kaan Ozkara; Antonious M Girgis; Deepesh Data; and Suhas Diggavi. A Generative Framework for Personalized Learning and Estimation: Theory, Algorithms, and Privacy. in International Conference on Learning Representations (ICLR). 2023.

2. One dataset seems not enough to demonstrate the scalability and generalization of the proposed framework.

We analyze three datasets: 1) the PV dataset for regression, 2) a polynomial synthetic dataset for regression, and 3) the FEMNIST dataset for classification. For the PV data, we experiment with diverse levels of heterogeneity. In both PV and FEMNIST datasets, we manipulate the number of samples per client. Consequently, we evaluate the performance of PAC-PFL across seven variants of these datasets and present the findings in Tables 2-6. Section 6 briefly introduces the datasets employed, and Appendix 8.5 provides additional details. In the updated version of the paper, we have provided a list of all datasets at the beginning of the experiments section to ensure clarity regarding the datasets we utilize.

With respect to scalability, our experiment on the FEMNIST dataset with around $500$ samples per client confirms that PAC-PFL scales reasonably well to realistically large datasets.

To enhance the richness of our experiments, we have introduced a new experiment utilizing the EMNIST dataset and the Dirichlet partitioning technique for introducing heterogeneity. The accuracy results for our method and some of the baselines are as follows: PAC-PFL ($87.1$%), FedAvg ($82.6$%), pFedBayes ($79$%), and Vanilla ($71$%). These results highlight the superior performance of PAC-PFL in handling the high heterogeneity introduced by the Dirichlet partitioning method. If the paper is accepted, we will dedicate a subsection to discussing the EMNIST dataset, accompanied by a comprehensive table showcasing the performance metrics of all baseline models. For detailed information about this experiment and its results, kindly refer to our response to question 3.2 from reviewer D4J2.

By incorporating datasets of distinct natures, i.e., regression and classification, and by systematically varying the characteristics of each dataset, such as the heterogeneity level and the size of the datasets, we believe that our experiments effectively demonstrate the capabilities of our framework.

3. The theoretical analysis is pretty solid. However, the experiments are not convincing and strong enough in contrast.

To ensure that our experiments effectively validate and showcase the proposed theory, we conducted simulations covering various aspects and features of the theoretical framework. To showcase the personalization capabilities, we utilized the highly heterogeneous and bi-modal PV-EW dataset. To highlight performance in scenarios with limited data, we varied the dataset size for both PV and FEMNIST datasets. To demonstrate the flexibility and generality of our method, we applied it to Gaussian Process regression and Bayesian Neural Network classification. By incorporating the new experiment on the EMNIST dataset, we substantiated the previously mentioned claims using a widely utilized benchmark.

To improve clarity regarding the connection between theory and experiments, we are now labeling the considered theoretical challenges as (c1-c4) in the introduction. These labels are utilized to illustrate how each experiment aligns with these theoretical challenges. These updates have been implemented in the revised draft. We thank you for the feedback. We believe that the experiments offer substantial support for the theory, but If you observe any other methodological results that seem underrepresented in the experiments, please kindly bring them to our attention.

Dear reviewer D4J2,

We appreciate your valuable feedback and thoughtful comments. Below, please find our answers to all your questions and concerns. We have submitted a revised version of the paper, encompassing the majority of the reviewers' comments. The main changes in the text appear in blue fonts. Should the paper be accepted, a comprehensive updated version will subsequently be submitted, addressing any outstanding comments. These updates, while relatively straightforward, were not included during the rebuttal phase due to timing constraints and are highlighted within specific responses.

1. The reason for using two samples per client remains unclear.

The rationale behind considering two distinct datasets for each client was explained in the fourth paragraph of the introduction and is detailed below. A selected client $i$ computes the global model’s update using locally available data denoted as $\mathcal{S}_{i}$. This is the first dataset considered for each client and aligns with the conventional definition of a client dataset in FL.

In practical scenarios, it is desirable to conduct global model updates at a lower frequency while collecting data and executing model personalization more frequently. For example, in the Photovoltaic (PV) panels experiment, it is likely that households download the updated shared model from the server and transmit model updates back on a weekly basis while measuring new data and personalizing their models on a daily basis. Our approach accommodates different time scales for updating the shared model and for personalization.

In between two consecutive updates of the global model, a client $i$ might collect a set of new data samples, denoted by $\tilde{\mathcal{S_{i}}}$. Although this data has not been utilized in training the shared model (i.e., the hyper-posterior), it can be leveraged by client $i$ for personalization. Considering an empty $\tilde{\mathcal{S_{i}}}$ set for all clients allows reverting to a typical FL framework. To the best of our knowledge, our work is the first to consider this second type of dataset, accounting for collecting new data over time. Another motivation for defining two datasets lies in providing a unified notation for existing clients and those joining the system later referred to as new clients. Given that a new client $\iota$ has never communicated with the server, $\mathcal{S_{\iota}}$ is an empty set. Conversely, $\tilde{\mathcal{S_{\iota}}}$ denotes the dataset that this new client has and can utilize for personalization.

In summary, by defining two datasets for each client, we formulate an FL framework capable of accommodating data samples used for personalization but excluded from federated training. These excluded samples may be collected by existing clients after sending model updates to the server or by new clients who have not yet transmitted updates to the server. If the paper is accepted we will include a more detailed explanation for using two different datasets.

2. The computation complexity of the algorithm can be very high. The communication overhead is increased from one to $k$ which can be unbearable for large models. It would be useful to see an ablation study on the choice of $k$.

The computational complexity of PAC-PFL is discussed in Appendix 8.3.7. PAC-PFL leverages the closed-form solution for the optimal posterior, effectively reducing the computations performed by each client. Consequently, the primary computational load in each iteration lies in calculating or approximating the log marginal likelihood (LML) given each SVGD particle. Given that computing the LML is an inevitable step in most hyper-parameter selection methods for probabilistic models, the increase in the computational cost is due to considering multiple particles, growing linearly in the number of particles, $k$. Similarly, as you mentioned, the communication overhead also scales linearly with $k$.

Notably, it has been demonstrated that SVGD achieves satisfactory performance with a relatively small number of particles (Liu and Wang, 2016). In our experiments, we employ 4 SVGD particles for the PV dataset and 3 particles for the FEMNIST dataset. If the paper is accepted, we will include a precise ablation study on the choice of the number of particles, $k$.

Given that both the computational complexity and communication overhead scale linearly with the number of SVGD particles and considering the obtained satisfactory performance with a modest number of particles, the overall complexity remains manageable. Our experiment on the FEMNIST dataset with around $500$ samples per client confirms that PAC-PFL can scale to realistically large datasets. Further reducing the computational complexity is identified as a potential avenue for future research in Section 7.

3.1. PAC-PFL even outperforms the Pooled method in Tables 2 & 3, which is surprising as Pooled is essentially an oracle. Additional explanation would be helpful. Also, the Pooled method is missing for FMNIST.

The Pooled GP method, as introduced in Section 6, involves training a single shared model on a dataset that concatenates data from all clients in a data-centric manner. Due to the absence of personalization in this approach, its performance is expected to be low, especially for heterogeneous clients. The explanation has been included in Section 6 of the updated paper within the paragraph introducing the baselines (please see the updated version of the paper). We thank you for highlighting this issue.

We acknowledge that the Pooled BNN results for the FEMNIST dataset were missing. The accuracy of the Pooled BNN model is $89.4$% when each client has $20$ training samples and $94.3$% when each client uses all the available data. The respective calibration errors are $0.106$ and $0.073$. These values have been added to Table 6 (please refer to the updated version of the paper). We thank you for bringing this to our attention.

3.2. It is common to use Dirichlet partition (Marfoq et al., 2021, Wang et al., 2020) for other image datasets to simulate heterogeneous clients.

Dirichlet partitioning is a method for artificially constructing a federated dataset with heterogeneous clients by partitioning a central and non-federated dataset. Each of these dataset partitions is then treated as the dataset for an individual client. This technique is useful when the central dataset lacks identifiers specifying the ownership of each sample by individual clients. For instance, the EMNIST dataset, a collection of hand-written characters, does not specify which writer produced each character.

In contrast, the FEMNIST dataset provides information about the writer for each hand-written character, facilitating a natural partitioning of the data among clients. Consequently, there is no need for artificial data division to create client datasets, as one can directly use the data from each writer as an individual client's data. In this scenario, heterogeneity arises naturally from the inherent differences in handwriting styles among various writers. This explanation aligns with the experimental setup outlined in (Marfoq et al., 2021), where Dirichlet partitioning is employed for the EMNIST dataset, and natural partitioning by writers is utilized for the FEMNIST dataset.

In response to your and reviewer QnoR's comments, we have performed a new experiment using the EMNIST dataset with Dirichlet partitioning. Following (Marfoq et al., 2021), we subsample $10$% of the EMNIST dataset and apply Dirichlet partitioning with the parameter $\alpha=0.4$ to synthesize heterogeneous clients. The objective of this experiment is to conduct a $64$-way classification task, recognizing hand-written characters that include both lower case and upper case English alphabet letters, as well as digits.

The accuracy results for our method and some of the baselines are as follows: PAC-PFL ($87.1$%), FedAvg ($82.6$%), pFedBayes ($79$%), and Vanilla ($71$%). For comparison, (Marfoq et al., 2021) reported an accuracy of $83.5$% on the same dataset. These results highlight the superior performance of PAC-PFL in handling the high heterogeneity introduced by the Dirichlet partitioning method. If the paper is accepted, we will dedicate a subsection to discussing the EMNIST dataset, accompanied by a comprehensive table showcasing the performance metrics of all baseline models.

(Liu and Wang, 2016) Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose Bayesian inference algorithm. In Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016.