Paper: Bayesian Coreset Optimization for Personalized Federated Learning (ICLR-2024)

Reviewer comment (Evaluation):
The major weakness is the lack of convergence comparison in the empirical part. One of the major concerns in FL is the communication cost. Thus the number of iteration rounds is crucial in FL. The reviewer suggests not only including the comparison of the final accuracy under (maybe different levels, not only 50%) of sample complexity, but also including the convergence speed, i.e., the communication cost comparison.

Author response:
We appreciate the reviewer's attention to the critical aspect of convergence comparison and its relevance to the communication cost in Federated Learning (FL). While we may have already provided in our theoretical results the convergence rate achieved in the case of our method which is in order of logarithmic bounds of the coreset size $k$ and thereby much faster than full data training,  we also get better convergence as compared to random sampling (as indicated in Fig 3). In addition to these, we update on our empirical results to include different subset sample sizes along with their communication cost (total no. of communication rounds) as follows.


In Appendix Section 10.2 we are including new results to show case the total number of communication rounds (where under each setting where we are considering different coreset subsample sizes k=50,30,15,10) which is indicative of the convergence of our proposed method. Empirical results shows that with significant lesser coreset subsample sizes we achieve faster convergence with less number of communication rounds , but however with lesser accuracy than that if we used a larger coreset size or full data training which is as expected. 



**Table: Comparative results of test accuracies across different coreset sample complexity**

| Method (Percentage = sampling fraction) | MNIST                       |                             | FashionMNIST                    |                             | CIFAR                       |                             |
|-----------------------------------------|-------------------------------|-----------------------------|-----------------------------------|-----------------------------|-----------------------------|-----------------------------|
|                                         | Test Accuracy                | Communication Rounds       | Test Accuracy                    | Communication Rounds       | Test Accuracy                | Communication Rounds       |
|-----------------------------------------|-------------------------------|-----------------------------|-----------------------------------|-----------------------------|-----------------------------|-----------------------------|
| $pFedBayes$ (Full)                      | 98.79                         | 194                         | 93.01                             | 215                         | 83.46                       | 266                         |
| $pFedBayes$ with random sampling (50%)                      | 80.2                          | 135                         | 87.12                             | 172                         | 48.31                       | 183                         |
| $Our Method$ (k = 50%)                  | 92.48                         | 98                          | 89.55                             | 93                          | 69.66                       | 112                         |
| $Our Method$ (k = 30%)                  | 90.17                         | 84                          | 88.16                             | 72                          | 59.12                       | 70                          |
| $Our Method$ (k = 15%)                  | 88.75                         | 62                          | 85.15                             | 38                          | 55.66                       | 32                          |
| $Our Method$ (k = 10%)                  | 85.43                         | 32                          | 82.64                             | 24                          | 48.25                       | 16                          |

Reviewer comment (Experiments):
How expensive it is to calculate the coreset samples/weights? Is there any empirical runtime results?

Author response:
Our approach to computing coreset samples/weights is exceptionally cost-effective, thanks to the utilization of accelerated iterative hard thresholding in the coreset—an incredibly efficient method for rapidly determining coreset weights. Although we haven't conducted explicit runtime experiments for coreset sample weights, we've assessed the overhead cost in terms of computation, revealing an average of only 0.001 seconds per round with a maximum of 10 iterations for coreset computation (which is ideal due to early convergance in coreset weights)  per round. Over 200 rounds, this results in an additional overhead cost of just 0.2 seconds compared to random sampling. 

We look forward to suggestions regarding details to any additional specific runtime experiments that the reviewer would like us to try out.

Reviewer comment (Writing):
How is \hat{\pi} defined in Eq. (3) and (4)?

Author response:
We have defined $\hat{\pi}$ as the weighting distribution that has the same support under an embedding Hilbert space as the true posterior $\pi$, for which we consider an L2 norm as the distance metric between the client's overall data likelihood and its corresponding coreset weighted data likelihood.

Reviewer comment (Writing):
Some typo: first sentence in section 3.2 is incomplete. Different places for \hat notation in q^i(\theta, w), on q, or q^i, or q^i(theta, w).

Author response:
Thank you for your careful inspection. We have identified this and similar places if any and have fixed the same in our revised version. Additionally regarding the notation for hat notation, $q^i(\theta, w)$ we have chosen to keep it consistent to $\hat{q^i}$. We have identified the sections where there is this notation inconsistency like in Section 4 and Section 6 of the main paper and have fixed them in our final revision. There are currently some portions in the supplementary which we are converting to a consistent notation for the hat operator and which will updated in our upcoming rebuttal revision.

Reviewer comment (Writing):
The paper's content is a bit bloated, and the use of notations can be messy. For instance, sections 3.2 and 4 could be condensed to make them more concise. Additionally, there is potential to simplify the formulaic aspect.

Author response:
We thank the reviewer for their valuable feedback and insights on the paper's content and notations. In response, we have streamlined and condensed sections 3.2 and 4 to enhance conciseness without compromising the clarity of our presentation, which has been reflected in our updated rebuttal version. We look forward to the reviewer's suggestions in understanding which other specific portions can be further simplified to improve upon our current draft.

Reviewer comment (Novelty):
It would be beneficial if the author could emphasize their novel contribution, distinguishing it from the techniques previously proposed by others. Currently, these ideas seem to be mixed within the intricate details of the interpretations.

Author response:
Our **primary objective** in this proposal is to **enhance the efficiency in Federated Learning training via utilizing a fraction of the original data at each client** while maintaining **near-optimal accuracy**, along with **minimally possible communication cost**. Our proposed theoretical findings showcase that **model trained on a client's coreset sample size weighted data show faster convergence rate** than in the case when trained fully on the entire training data. 

Both Bayesian Coreset Optimization and Federated Learning Optimization problems are **two different orthogonal and independent problems**. Considering their both distinct challenges, we aim to infuse the two optimization framework in one single unified cohesive framework such that the coreset optimization problem can **strategically pick data points for each client** in a **model dependent manner**, i.e. without facing too much penalty in the model training under Federated Learning setup. As of now, to the best of our knowledge, this specific problem **remains largely unexplored**. Remarkably, there is a significant gap in the existing research landscape concerning subset selection or coreset-based strategies within the context of Federated Learning, with a conspicuous absence of theoretical frameworks in this domain.

Crucially, **our work is unique in addressing the absence of research on subset selection or coreset-based strategies within Federated Learning setup** and theoretical contributions along with empirical results in terms of near optimal performance under data scarcity setting accompany our contributions.

Lastly we also **showcase its applicability in medical domain scenario** via medical datasets where **data scarcity remains a huge problem** and how our proposal can alleviate the same.

Reviewer comment (Novelty):
The overall architecture, as well as certain smaller techniques and theoretical analysis methods, seem to be largely derived from previous work.

Author response:
We thank for the reviewer's feedback in this aspect. However, we respectfully assert a differing viewpoint. While building upon prior research to establish a robust foundation, our **contribution is distinctive in its innovative synthesis and adaptation of Bayesian Coresets and Personalized Federated Learning frameworks**. The fusion of these domains creates a novel architecture of its own as it can be seen in Algorithm 1 with unique optimization objectives, thereby further leading to new theoretical formulations and findings, as evidenced in our Theoretical Analysis and extensive novel formulations in the Supplementary (Section 9 Lemma 1, Proposition 1, Proposition 3, Proposition 4. ).  We further refine and extend these objectives to enhance performance and efficiency. We firmly believe that the synergy of existing knowledge with our novel contributions significantly strengthens the overall impact and applicability of our work. We aim to further clarify certain aspects of the above thread with specific instances as we go forward.

Reviewer comment (Novelty):
The contribution on the coreset construction is limited. Although the authors introduce a new coreset objective, they do not provide sufficient NEW optimization techniques for the new objective. I could only identify some techniques borrowed from previous work.

Author response:
We appreciate the reviewer's feedback. It is essential to emphasize that our primary contribution lies not in introducing new optimization techniques for the Bayesian Coreset construction but instead, our innovation centers around **proposing a comprehensive algorithm and corresponding optimization strategies for a common unifying framework** that governs two distinct optimization problems at hand. This approach leads to a **model-dependent coreset sampling technique** (in this case the model being trained under a personalized federated learning setup) i.e. along with the model training for federated learning, the coreset weights are also learnt dynamically for each client, which is a novel aspect of our contribution. While drawing inspiration from previous work is a common practice, our work's strength lies in the integration and adaptation of these techniques to address the challenges posed by our unique framework. In particular for bayesian coreset optimization problem, we **propose a novel objective loss in Eq 9/8** on top of which we apply Accelerated Iterative Hard Thresholding. It's noteworthy that while AIHT is an efficient optimization strategy in its own right, our emphasis is not on its intrinsic properties. Rather our **proposed optimization objective stands out for its considerable complexity** and to address the same, we **employ optimization strategies tailored to its nuances**, such as the use of AIHT. Detailed insights into this choice are provided in the Supplementary section proofs, affirming our confidence in the appropriateness of this approach.

Reviewer comment (Novelty):
In my opinion, the primary contribution of this paper is the modified objective (eq. 9) tailored to personalized federated learning. However, the advantages of this modified objective are not adequately elucidated in the current presentation.

Author response:
We appreciate your recognition of the main contribution being the modified objective function (Eq. 9) tailored for personalized federated learning. Your feedback on the need for a more thorough elucidation of the advantages is duly noted. In Section 3.2, we delve into the rationale and motivation behind our proposed approach, offering a more comprehensive understanding of the modified objective's benefits. Additionally, we have conducted new experiments and updated our explanations in response to Question 2 under **Questions**, specifically addressing the advantages of the modified objective. This enhancement aims to provide a clearer and more detailed exposition of the strengths associated with our proposed contribution.

Reviewer comment (Writing):
Many other typos, e.g., missing equation references and confusing sentence like “For the first term in Equation 1, the authors we use a minibatch stochastic gradient descent …”

Author response:
**minor typos**: We have identified the same and currently fixing the same for our next rebuttal version. We will update shortly.

Reviewer comment (Experiments):
What is the benefits to apply coreset in the personalized federated learning? I think one of the most important is that it can reduce the communication complexity. It would be valuable to investigate and quantify the extent to which the coreset approach reduces communication complexity in the specific optimization task addressed in this work. This can be done theoretically, by providing a complexity formula, and practically, by presenting numerical results from experiments that show the reduction in communication complexity achieved.

Author response:
We value the reviewer's input on comprehending the computational complexity and the effectiveness of our approach in this context. Consequently, we perform an analysis of communication complexity to assess the total number of communication rounds required for full convergence across various coreset subsample sizes. We present our results as follows (Additionally we have included these results under Supplementary section 10.2)


| Method (Percentage = sampling fraction) | MNIST                       |                             | FashionMNIST                    |                             | CIFAR                       |                             |
|-----------------------------------------|-------------------------------|-----------------------------|-----------------------------------|-----------------------------|-----------------------------|-----------------------------|
|                                         | Test Accuracy                | Communication Rounds       | Test Accuracy                    | Communication Rounds       | Test Accuracy                | Communication Rounds       |
|-----------------------------------------|-------------------------------|-----------------------------|-----------------------------------|-----------------------------|-----------------------------|-----------------------------|
| $pFedBayes$ (Full)                      | 98.79                         | 194                         | 93.01                             | 215                         | 83.46                       | 266                         |
| $pFedBayes$ with random sampling (50%)                      | 80.2                          | 135                         | 87.12                             | 172                         | 48.31                       | 183                         |
| $Our Method$ (k = 50%)                  | 92.48                         | 98                          | 89.55                             | 93                          | 69.66                       | 112                         |
| $Our Method$ (k = 30%)                  | 90.17                         | 84                          | 88.16                             | 72                          | 59.12                       | 70                          |
| $Our Method$ (k = 15%)                  | 88.75                         | 62                          | 85.15                             | 38                          | 55.66                       | 32                          |
| $Our Method$ (k = 10%)                  | 85.43                         | 32                          | 82.64                             | 24                          | 48.25                       | 16                          |

Reviewer comment (Theory):
The intuition behind the new objective in eq. 9 is not very persuasive. If you could compute a coreset with a sufficiently small loss as defined in eq. 3, it is unnecessary to add the term representing the “distance” between $\hat{q}^i(\theta, w)$ and $\hat{q}^i(\theta)$ since $\hat{q}^i(\theta, w)$ and $\hat{q}^i(\theta)$ will lead to closed losses; on the other hand, if you couldn’t make it under the constraint $\| w \|_0 \leq k$, which means there is no such small coreset with ideal error, the coreset method could not work well. It would be beneficial to clarify the merits of the new objective, such as its robustness or any other advantages it offers. Experiments that demonstrate the effectiveness of the new objective would greatly strengthen the argument.

Author response:
**Response**
We appreciate your thoughtful engagement with our work and would like to address the concerns regarding the intuition of the proposed objective in Equation 9. Your point about the potential redundancy of the first term in Eq 9. i.e. the KL divergence term, given that we are already considering the second term i.e. the coreset loss term, is a valid claim. However, we would like to offer a different viewpoint regarding the same. In addition to our discussions in Section 3.2, we provide further clarifications as to our motivation behind including the first term in Eq 9, along with corresponding empirical observations to justify the same. Note the second term captures the deviation of the data likelihood over the entire client's data w.r.t that of the coreset weighted client data. When the second term loss is small under specific sparsity constraints, it leads to the conclusion that a suitable $\boldsymbol{w}$ has been found such that coreset weighted client's data likelihood mimics to a good extent to that of the original client's data likelihood. The above of choice of $\boldsymbol{w}$ can surely be then utilised to achieve a corresponding solution to the minimization problem of client side objective: Eq 6. (which is say $\hat{q^{i}}(\theta, w)$. \textbf{However, this optimal $\hat{q^{i}}(\theta, w)$ distribution (i.e. local distribution of the coreset weighted client) can be any arbitrary distribution acting as a minima to Eq 6. (as there might be many such local minimas for the same), which further leads to a different optimal variational solution for the global problem in Eq 2. (say $\boldsymbol{z^*(\boldsymbol{\theta},w)}$ ) , which can be quite different than the optimal variational solution for the global setting in the case when the entire client's data is considered (say $\boldsymbol{z^*(\boldsymbol{\theta})}$), thereby not leading to an optimal solution at the end.

This was also initially observed empirically when we just tried out with the coreset loss as our main objective function. In Supplementary section 10.1 we do a study to check why the KL divergence term (more specifically forcing the local client's distribution under a coreset setting mimic that of the local client's distribution under a full data training setting) is beneficial.

Based on our experimental analysis, we observe that exclusively incorporating the coreset loss optimization may lead to early saturation, suggesting a potential susceptibility to getting trapped in local minima. However, introducing the KL Divergence loss and enforcing similarity between the coreset-weighted local distribution of the client and the normal local distribution yields improved stability in the training loss and enhances convergence.

Corresponding plots : https://pasteboard.co/3VGSZG0B4ab1.png (MNIST)
                                  https://pasteboard.co/4pXBqEIhy5Ik.png (CIFAR)

Reviewer comment (Theory):
Does the modifications of eq. 6 consist of the following two parts: i) use the weighted likelihood; ii) replace prior distribution with global distribution. I am not sure for that.

Author response:
**Response** : Yes correct, the modifications of eq. 6 consist of the above parts

Reviewer comment (Experiments):
Is there any strategy for choosing the value of k in practice?

Author response:
**Response** : The consideration depends on the specific application setting we are looking at in terms of resource/data constraints and the level of sacrifice in terms of accuracy acceptable for the given context. While we cannot directly quantify the accuracy in relation to k, for a particular k (since a lot of other factors weigh in like class-wise distribution and others), the goal here is to achieve optimal accuracy that approaches the original accuracy with considerable gaps. The determination is influenced by the trade-offs we are willing to make.

Reviewer comment (Novelty):
I am not sure what is the challenge in incorporating the Bayesian coreset framework in federated learning setting. It would be better to explain clearly why this is a significant contribution. Both the algorithm and proof techniques appear to be heavily inspired from Zhang 2022b. The only modification seems to be use of Bayesian coresets.

Author response:
One of the **core problems that we try to address** here is how to **achieve efficiency in training in a Federated Learning setup** **using only a fraction of the original data at each client** yet achieving **near optimal performance in terms of accuracy** (theoretical results on our logarithm bounded convergence rate using coreset indicates better convergence than full training while empirical results across a diverse set of baselines including diversity based subset selection techniques show promising results in terms of accuracy). Both the Bayesian Coreset optimization framework and Personalized Federated Learning framework  which in itself involves a  bilevel optimization problem are two separate entirely different orthogonal problems. 

Our goal in this proposal is thus to **bridge the two towards a common optimization framework** whereby we can utilise a Bayesian coreset framework for strategic selection of data points at each client level and at the same time we do not need to pay for a bigger penalty while solving the Federated Learning optimization problem. Both these problems represent independent optimization challenges that need to be addressed concurrently. The difficulty arises when striving to learn the coreset in an optimized manner that minimizes the overall penalty incurred while training within the Federated Learning setting. Hence our work focuses on seamlessly **interleaving these two optimization frameworks** into a **single, cohesive framework**, introducing **novel optimization formulations as evident in Eq 8 and 7** (updated equation numbers due to manuscript revision), along with corresponding strong theoretical guarantees. Importantly, our *contribution addresses a gap in the existing literature*. To the best of our knowledge, this specific problem has not been thoroughly explored. There is a notable absence of research on subset selection or coreset-based strategies within a Federated Learning context, let alone the formulation of theoretical frameworks in this domain. 


Additionally, as an application setting we also showcase the usefulness and applicability of our method on **real-world medical datasets where in many cases there is a severe dearth of medical data**.
To address the challenges in infusing the two optimization frameworks together and propose new findings on top of it, we draw upon some of the proof techniques from papers like Polson
& Roˇckov ́a (2018), and others, one of them being Zhang 2022b, particularly due to the similarities observed in Eq. 1. We utilize certain existing lemmas and theorems, on top of which we provide novel theoretical formulations and new results as seen in theorem 1 and theorem 2, and their accompanied proofs in Supplementary (Section 9 Lemma 1, Proposition 1, Proposition 3, Proposition 4. ) which are our primary contributions.

[1] Nicholas G Polson and Veronika Roˇckov ́a. Posterior concentration for sparse deep learning.

Reviewer comment (Writing):
There are minor grammatical errors. Please do a grammar check.

Author response:
We thank the reviewer for their careful inspection. We have made a thorough proofread pass and have identified the specific areas and have fixed them in our recent revision.

Reviewer comment (Writing):
Why the prior $\pi$ in equation 1 is replaced by $\mathbf{z}$ in eq.6 - the modified client-side objective. Please clarify.

Author response:
Thank you for your inquiry regarding the substitution of the prior in Equation 1 with 'z' in Equation 6. In Equation 1, we initially introduce the concept of a prior distribution when formalizing the client-side objective. For the sake of clarity and ease of understanding, we denote this prior probability distribution as $\pi$ as it represents a more nuanced notation to indicate probability and also as it encompasses additional terms such as $\pi(\theta|D)$. In Equation 6, we streamline and simplify the representation while retaining the essence of the prior distribution in a more concise form using *z*. We have also adequately defined after the equation to resolve any confusion for the same.

Reviewer comment (Experiments):
The subsample size is 50%. Is it not quite large? Does it give significant computational time benefits when compared with full data? Other than figure 3, there are no experiments mentioning computational efficiency.

Author response:
We do agree that taking subsample size = 50% can be considered somewhat large, but at the same time dropping too many data points may also result in degradation in performance (accuracy) irrespective of what sampling strategies we use. Hence we thought 50% to be a respectful threshold where we can measure the effectiveness of our proposed approach for maintaining near-optimal performance as compared to other baselines. Keeping in consideration the importance of studying across different thresholds we have included some of the baseline comparisons on less percentage data as follows.

In Appendix Section 10.2 we are including new results to show case the total number of communication rounds (per setting where we consider different coreset subsample sizes k=50,30,15,10) which is indicative of the convergence of our proposed method. Empirical results shows that with significant lesser coreset subsample sizes we achieve faster convergence with less number of communication rounds but with lesser accuracy than that if we used a larger coreset size or full data training. 
**Table: Comparative results of test accuracies across different coreset sample complexity**

| Method (Percentage = sampling fraction) | MNIST                       |                             | FashionMNIST                    |                             | CIFAR                       |                             |
|-----------------------------------------|-------------------------------|-----------------------------|-----------------------------------|-----------------------------|-----------------------------|-----------------------------|
|                                         | Test Accuracy                | Communication Rounds       | Test Accuracy                    | Communication Rounds       | Test Accuracy                | Communication Rounds       |
|-----------------------------------------|-------------------------------|-----------------------------|-----------------------------------|-----------------------------|-----------------------------|-----------------------------|
| $pFedBayes$ (Full)                      | 98.79                         | 194                         | 93.01                             | 215                         | 83.46                       | 266                         |
| $pFedBayes$ with random sampling (50%)                      | 80.2                          | 135                         | 87.12                             | 172                         | 48.31                       | 183                         |
| $Our Method$ (k = 50%)                  | 92.48                         | 98                          | 89.55                             | 93                          | 69.66                       | 112                         |
| $Our Method$ (k = 30%)                  | 90.17                         | 84                          | 88.16                             | 72                          | 59.12                       | 70                          |
| $Our Method$ (k = 15%)                  | 88.75                         | 62                          | 85.15                             | 38                          | 55.66                       | 32                          |
| $Our Method$ (k = 10%)                  | 85.43                         | 32                          | 82.64                             | 24                          | 48.25                       | 16                          |

Reviewer comment (Writing):
Not a question but a suggestion. Algorithm 1 is not easy to follow for anyone unfamiliar with existing work or similar algorithms. How exactly is the coreset getting constructed? It would be good to give a high-level description of the same.

Author response:
In Algorithm 1, our approach involves evaluating the client side objective function both the full data-based likelihood and the coreset weighted likelihood for each client. Subsequently, we calculate the objective function using Equation 8 ( now changed Eq number due to revised manuscript). To compute the coreset samples, we employ the highly effective Accelerated IHT algorithm atop this objective function. Supplementary Section 10.5 provides comprehensive details on the Accelerated IHT algorithm for a deeper understanding of its application in our methodology.