A Bayesian Framework for Clustered Federated Learning

We thank the reviewer and we reply to the points raised by the reviewer.

• The notation complexity is due to the indexing of the multiple association hypotheses, being random finite sets themselves, which is at the core of the conceptual solution of the association problem brought by clustering. It would be helpful if the reviewer could point out notation issues so we could address them. Notice that there is a summary table with notation in the appendix and that we are revising it to make sure relevant variables are defined.
• At every communication round we use the Laplace approximation approach as discussed in "Baseline models and system settings" in Section 5.
• We agree with the reviewer that the discussion presented in the related work section could be improved. We are working on improving the discussion for the revised manuscript.
• We strongly disagree with the reviewer on this point. We formulated the problem of clustered federated learning as a Bayesian data association problem between clients and model parameter distribution. The theoretical approach presented in Section 3 is completely novel in the FL literature: section 3.1 presents the conceptual formulation and solution to compute the joint posterior of model parameters and data association; section 3.2 extends 3.1 to recursive Bayesian updates and associations, which are implemented as FL iterations and yield to an optimal yet intractable solution as pointed out in section 3.3. Furthermore, the derivations are far from trivial, they are rigorous in the sense that they rely on consolidated Random Finite Set and Bayesian theories with few assumptions, no approximations or heuristics. The proposed approach is a unified framework since the theory proposed in Section 3 is general and can lead to multiple solutions as those in Section 4. The proposed conceptual solution, however, leads to a high number of data associations and, consequently, to high computational complexity if applied naively. Nevertheless, the insights provided by the theory laid out in Section 3 were used to propose the approximations discussed in Section 4 which successfully aim at reducing the complexity through novel algorithmic solutions. We modified the title of Section 3 as "Optimal Bayesian solution for clustered federated learning" to make it clear that the material there is already part of the novel contribution of this paper, which we agree was not clear earlier when titled "Problem formulation". Overall, we believe that the proposed approach touches all the points we claimed. We hope that the reviewer can reconsider such a strong claim that the paper fails to present any novel theoretical contribution.
• We disagree, see our previous answer.
• The reviewer is mistaken. Section 3 and 4.1 present different pieces of the proposed BCFL: while section 3 presents the conceptual solution where all association hypotheses are kept, section 4.1 presents the basics to weight different hypotheses and design approximate solutions to the computational challenge. The solution of equation (16) leads to the set of data associations that maximizes the sum, over all $K$ clusters and all clients in each cluster, of the $\log$ unnormalized marginal posterior of data association weights $\tilde{\pi}^{\theta^{ij}_t|h}$, where $h\triangleq \theta\_{1:t-1}$ indicates the history of past associations as discussed below equation (11). In the case where Equation (9) is approximated using the expected value of the model parameters then $\tilde{\pi}^{\theta^{ij}\_t|h} = p(\mathcal{D}^j|\omega^i\_{t-1})$. Nevertheless, only in the greedy approach discussed in Section 4.2 the selection is made solely on the $\log p(D^j|\omega^i\_{t-1})$. We hope that the reviewer can revisit this point and change the overall evaluation of the manuscript.
• We will add these references to the related work section discussing it connects to BCFL. Note that a FedSoft is already cited in our work. In our paper, we focused on the clustered FL problem as stated in (Ghosh et al., 2020; Mansour et al., 2020), where the objective is to learn multiple models under the assumption that data from each client comes from a specific distribution.
  Our approach focuses in clustering the different clients using a data association Bayesian framework that provides a joint posterior distribution of client associations and model parameters. Opposed to existing methods, we consider multiple data association hypotheses that enable the uncertainty quantification of client-to-cluster associations, leading to noticeable performance improvement. In EM-based FL approaches the assumption is that the data at each client belongs to a mixture of distributions. Then, FL approaches are used to independently train multiple models under this assumption, which does not explore different association hypotheses. Furthermore, FedEM still relies on frequentist estimation strategies which contrasts with our Bayesian framework.

We thank the reviewer for the appreciation of the strengths of our work and the feedback provided. Below we reply point by point to the weaknesses perceived by the reviewer.

• There are three comments within the first item. Namely:

1. It is worth clarifying that downloading the model for all past hypotheses is only the case for the conceptual solution discussed in Section 3 of the paper. However, the approximate solutions proposed in Section 4 are aimed to reduce those communication needs and thus to improve the practicality of BCFL. Other things to consider are that i) the weight associations variable sent along the model weights are just a few scalar values, thus requiring minimal extra communication cost; and ii) this paper tackles the additional communication cost by way of the proposed approaches that reduce the need to send an otherwise growing number of hypotheses.
2. We also conducted tests involving a larger number of clients and have presented additional examples with 40 clients in the appendix, we will emphasize this in the updated paper. In the main body, we included only a limited number of clients for the Digit5 and Amazon Review datasets to ensure the clarity of our cluster analysis results, as done in other related works. As depicted in Figure 4, the behavior of the clusters during training is distinctly observable.
3. Regarding the knowledge of the number of clusters $K$, we preliminary followed common practice in CFL works where this is often predetermined. We agree that estimating $K$ is of interest and it is actually on our future research plan, which our Bayesian framework can account for although not straightforwardly enough to be included in this paper so we instead decided to keep it fixed (although not necessarily correctly specified). To address this comment we updated the appendix with an experiment where the sensitivity to under- and over-estimating $K$ is discussed.

• We would like to clarify that the main assumptions in the paper are that given an hypothesis the model parameters per cluster are independent and that the local datasets are conditionally independent. The intuition behind the validity of those assumptions comes from the conditioning on the clustering hypothesis, whereby if we were to explore all of the hypotheses we would find that different clusters obey different models. We agree that this might not be always true, that is the reason behind explicitly stating it as an assumption. Doing so, we enable tractability of the problem, which otherwise would be extremely challenging. Finally, it is worth noting that other non-Bayesian CFL works implicitly make this assumption, although since those works are often not probabilistic this might not be as apparent as in our framework.
• Thanks for pointing out the typos and figure suggestions, they are fixed. As for the use of `optimality', we refer to the fact that the conceptual solution in Section 3 provides the full posterior distribution of the model, where all association hypotheses are explored. That is, BCFL is optimal in the Bayesian sense, where the objective is the computation of the joint posterior distribution of the model parameters and the clustering hypotheses. Optimality of the algorithms is thus referring to a Bayesian inference notion of optimality, which is now mentioned in the first paragraph of Section 3. The approximate solutions developed in Section 4 are therefore suboptimal, although they ultimately aim at approximating that same joint posterior while maintaining tractability and practicality of BCFL.

We thank the Reviewer for the appreciation of the great summary and strengths of our work by pointing out the contribution related to data association and Federated Learning.

1. Presentation:

We improved the notation readability for those sections, as well as added pseudo-code algorithmic description for both the conceptual solution (section 3) and the practical schemes (section 4). These are in Section B of the Appendix. We acknowledge the possibility of enhancing these sections for clarity. However, given the constraints of time, it is challenging to undertake substantial revisions in less than a day. Nevertheless, we are committed to addressing the reviewer's concerns to the best of our ability within the time available and in future iterations.

2. Related work:

Thanks for the pointers regarding how to improve the related work section. We improved it by adding more references and better motivating the contribution of this work within those works. Particularly, we discussed Bayesian frameworks within FL, as the reviewer suggested, which has indeed improved the positioning of this paper.

3. Method:

Thanks, as the reviewer points out, the 3 methods proposed in the work are approximated because we cannot account for all the associations (discussed in section 3.3). More precisely: BCFL-G only keeps the best association every time; BCFL-C chooses the $M$ best associations at a given time instant and merges them into a single posterior solution, thus only propagating one hypothesis to the next communication round; and BCFL-MH keeps the $M$ best associations onto the next recursively update. There might be other approaches to approximating the full posterior, which itself can lead to different BCFL algorithms. Indeed, the results for BCFL-MH-6 do not always improve those from BCFL-MH-3. There could be multiple reasons to that such as related to the dataset and its setting, the chosen associations, and the randomness of each trial. Ultimately, when taking a close look at the result, MH-3 or MH-6 are noticeably close to each other, which indicates that for those experiments going from $M=3$ to $M=6$ does not provide substantial performance improvement.

4. Experiment:

• All these datasets are real datasets. The usage of those datasets in the context of non-IID challenges is according to similar CFL works, which we compare to. We do not understand which questions regarding the motivation of CFL the reviewer is referring to. More details about the datasets are provided in the Appendix.
• Per the reviewer's suggestion, we included a discussion on the communication efficiency for each practical algorithm in Appendix B.

5. Minors:

(1) We provide a bit more discussion of the claim in the `Problem 1' description in the introduction section. The statement regarding other CFL methods is that the knowledge of each participating client is exploited by only one cluster during each round, which potentially results in an inefficient utilization of the local information that could contribute to the training of multiple clusters instead;

(2) We improved it and more details can be found in a newly added Appendix C, as well as the references therein;

(3) We implemented the warm-up method to see if it could help achieve faster convergence. And indeed, it did help. With the same number of communication rounds, it usually achieves better results. The reason behind this is that we provide a good initial estimate at the beginning, as shown in Figure 3. Warm-up methods typically have good initial accuracy and are used in other works.

Questions:

We improved the notation readability for those sections, as well as added pseudo-code algorithmic description for both the conceptual solution (section 3) and the practical schemes (section 4). These are in Section B of the Appendix. These pipeline descriptions should help improve the presentation of the work, as well as the information flow in the different methods (including the conceptual solution.

Regarding the option of simply presenting BCFL-MH, our main objective is to present the BCFL framework from which one can derive the three proposed solution and many others that can be explored in future research from us or other groups. The objective of this work is to present a new framework for the problem of clustered FL, rather than just presenting a detailed method. The three approaches are practical implementations, derived from the framework. There could be many other variations by pruning, merging, and combining them, which would expand the possibilities to make progress in the CFL field. The BCFL-MH approach yields better results because it considers more associations compared to the other approaches. However, it is not as efficient as others, as discussed in the communication section of the appendix. That is the reason why we presented the three alternatives, balancing communication needs and performance. Each application might have different requirements.

We thank the reviewer and we reply to the points raised by the reviewer.

Weakness:

1. We improved the notation readability for those sections, as well as added pseudo-code algorithmic description for both the conceptual solution (section 3) and the practical schemes (section 4). We acknowledge the possibility of enhancing these sections for clarity. However, given the constraints of time, it is challenging to undertake substantial revisions in less than a day. Nevertheless, we are committed to addressing the reviewer's concerns to the best of our ability within the time available and in future iterations.

2. Firstly, the methods we propose are not only for experimental comparison only, but our primary objective is to introduce a framework that can be adapted and extended to accommodate a variety of methods for future research. Secondly, WeCFL represents a non-trivial algorithm, and we do not claim that our methods will universally outperform it. However, as indicated by the results, BCFL-G/C achieve similar or improved performance compared to WeCFL across a diverse range of datasets and experimental configurations, with some additional results presented in the Appendix.

3. The main contribution is the Bayesian framework for CFL, the specific algorithms like BCFL-MH are options within it. While BCFL-MH seems to be a good solution performance-wise, the other BCFL-G/Ca approaches provide remarkable performance at lower communication/complexity burdens. We added a discussion on this in the revision, as well as a discussion of the communication and computation cost in Appendix, Section B.

4. The theoretical insights come from the formulation of the clustering FL problem as a probabilistic model which we address from a Bayesian perspective, dealing with the multiple data associations. This allowed us to propose different practical methods, solidly based on the Bayesian paradigm. We acknowledge the reviewer's concern regarding the absence of a formal guarantee in this part, or at least guarantees in the sense of those provided in frequentist-based methods. As a matter of fact, our proposed approximation BCFL-G has strong connections to the IFCA [ghosh2020efficient], when formulated in a Bayesian context, which provides those theoretical guarantees. Our paper concentrates on the Bayesian framework rather than on particular deterministic algorithms, which is why we did not initially include a convergence proof. However, we recognize the importance of such an analysis, and we are certainly open to investigating convergence guarantees in future research.

5. We added sensitivity analysis. Please refer to our appendix for results on the selection of $K$ and its justification.