On Bridging Generic and Personalized Federated Learning for Image Classification

Abstract: Federated learning is promising for its capability to collaboratively train models with multiple clients without accessing their data, but vulnerable when clients' data distributions diverge from each other. This divergence further leads to a dilemma: "Should we prioritize the learned model's generic performance (for future use at the server) or its personalized performance (for each client)?" These two, seemingly competing goals have divided the community to focus on one or the other, yet in this paper we show that it is possible to approach both at the same time. Concretely, we propose a novel federated learning framework that explicitly decouples a model's dual duties with two prediction tasks. On the one hand, we introduce a family of losses that are robust to non-identical class distributions, enabling clients to train a generic predictor with a consistent objective across them. On the other hand, we formulate the personalized predictor as a lightweight adaptive module that is learned to minimize each client's empirical risk on top of the generic predictor. With this two-loss, two-predictor framework which we name Federated Robust Decoupling (Fed-RoD), the learned model can simultaneously achieve state-of-the-art generic and personalized performance, essentially bridging the two tasks. 

Reviews and author responses:

FL trains models both on the server and clients. Such model training becomes problematic when the data distributions on each of the clients diverge from each other. The paper answers the question of prioritizing the generic model at the server or the personalized model on each client. The paper proposes to decouple a model’s dual duties with two prediction tasks.  ### Summary:

FL trains models both on the server and clients. Such model training becomes problematic when the data distributions on each of the clients diverge from each other. The paper answers the question of prioritizing the generic model at the server or the personalized model on each client. The paper proposes to decouple a model’s dual duties with two prediction tasks. 

### Pros:

- The paper is well-written, easy to understand and follow along.

- The paper is well motivated and driven by a real-problem in FL. 

- Considering model training on each of the clients as a class-imbalance problem and applying minimizations on the objectives considering the two losses in a multi-task setup is also interesting.

### Cons:
Following are the concerns that can be addressed for change in the final scores.

1. “the literature does not evaluate G-FL algorithms in the P-FL setup by their local models” What does this mean? The evaluations on global test sets or the local test sets? In either case, this sentence does not really make a lot of sense.

2. “taking average over model weights -- indeed acts like a regularizer” is also not a surprising finding, this indeed is general common sense to think that averaging is kind of a regularizer.

3. In fact look at the following paper about bagging (which of course is not an FL algorithm but a simple general ensemble technique of averaging multiple models) that suggests to have regularization effects due to model averaging.
    - Poggio, T., Rifkin, R., Mukherjee, S. and Rakhlin, A., 2002. Bagging regularizes. MASSACHUSETTS INST OF TECH CAMBRIDGE ARTIFICIAL INTELLIGENCE LAB.
    - Skurichina, M. and Duin, R.P., 1998. Bagging for linear classifiers. Pattern Recognition, 31(7), pp.909-930.

4. If one performs the P-FL evaluations on local test sets, at the beginning of each round with the newly communicated global model that might perform better on local test sets as opposed to weight for the local model to be fine-tuned and perform testing, not really surprising, even without regularization term because of the implicit reg effect due to averaging.

5. In P-FL, the local test sets used in the experiment are static because of the standard benchmarks used. However, in real-time, these test sets will change and there will be temporal effects that need to be considered. In such cases, building test sets in itself becomes a challenging problem. The whole paper starts with the motivation of closing the gap between G-FL vs P-FL, the gap arises due to the data shifts, for which the corresponding test sets need to be prepared. This issue is not addressed at all.

6. y_P = h^G+h^P, which is uniform weight for the P-FL part, how about you add weights and make them learable, something like this y_P = (1-\alpha)h^G + \alpha h^P and start with bigger \alpha as this function optimizes for the P-FL anyway, why not try something along these lines and see how that improves P-FL and effects the G-FL. Any of the findings in this direction are also useful since you already proposed this and have the setup.


### Post Rebuttal
Satisfied with the authors response, increasing the score to a clear accept.
 The idea of treating the non-IIDness on clients in FL as class imbalance is interesting, however, the paper can benefit from addressing the above raised concerns.

This paper focuses on a challenging problem in federated learning setting where data across different local workers are imbalanced, and most of the current work in the community are optimizing either generic performance or personalized performance. This paper proposes a new framework, Fed-RoD, which combines both global objective and local objective, to achieve that the trained model is suitable for both generic and personalized task. The main contributions are:
1. Focuses on an interesting problem which aims to increase both local and global performance in imbalanced federated learning setting;
2. Designed a two-loss, two-predictor framework to achieve the goal;
3. Based on the extensive experiments, the performance achieves state-of-the-art on both sides. Strengths:
1. The problem is clearly stated and the algorithm is well-organised, also easy to understand and follow;
2. The problem itself is interesting, non-iid and imbalance is a key challenge in federated learning and real-world setting. Also, increase both local and global performances should be of interest to the community;
3. The experiments are well designed and comprehensive;
4. The experimental results provided in the paper do achieve a state-of-the-art on both generic and personalized setting;
5. Related works are properly cited.
 
Weaknesses:
1. Since no code is provided, it is useful to include more training details, for example, training hyperparameters for readers to better utilize the proposed framework;
2. Also for the experiments, how many repeat trials did you run for a experiment? It is important to also include the error bar in the results;
3. Lack of theoretical guarantees on the proposed framework. This paper proposes a new framework aiming to increase generic and personalized performance simultaneously in the imbalanced federated learning setting, the proposed algorithm is simple yet efficient according to the experimental results. Theoretical analysis are somehow insufficient, and some minor fix or explanations are needed to better understand the experiments. 

 I thank the authors for the response. My concerns are well addressed, and considering the work has more applicable usage even without theoretical analysis, I am happy to change my score.

 Dear reviewers,

Thank you again for reading our rebuttal.

We have tried to address most if not all concerns raised by the reviewers. Please let us know if you have any further questions about our paper or rebuttal. If our responses have sufficiently addressed your questions/concerns, it will be great if you are willing to kindly reconsider your score. Thank you.

Best,

authors

 **4. P-FL evaluation on local test sets using global models vs. local models.** Thanks for the comment. We respectfully think there might be a misunderstanding. When we mentioned the implicit regularization effect in Section 3.3, we mean that performing model average before local training helps improve the fine-tuned local models’ personalized performance. However, it does NOT mean that the global model (at the beginning of each round) would perform well on local test sets. This is evidenced in the bottom-right subfigure (P-FL Acc plot) of Figure 1. The fine-tuned local models $\{w_m\}$ are in fact better than the global model $\bar w$ in the P-FL evaluation. This is because the local test sets are fairly different in their distributions and a single global model without personalization cannot perform well on all of them.

Another evidence is in Table 1. There, P-FL means testing on local test sets. Under each P-FL column, there are two sub-columns: GM means applying the global model; PM means applying the local/personalized models after local training. We can see that for FedAvg, PM outperforms GM on the P-FL accuracy. However, when we compare PM of FedAvg to PM of Ditto, where the former has the implicit regularization (Section 3.3) while the latter does not (Ditto imposes an explicit one as in Equation 4), we see that the implicit regularizer --- by model average before local training --- leads a to higher P-FL accuracy in general.

**5. On the design of local test sets.** We thank the reviewer for the insights and great suggestions. We want to first clarify that our work follows the current literature of P-FL to study “static” local test sets that have different data distributions across clients. Our motivation for bridging generic and personalized FL is to simultaneously learn a strong generic model for the global test set and strong personalized models for local test sets, especially when the clients’ data distributions are very different. We agree that the dynamically changing local test set is very interesting, and we will consider it in our future work. Potential datasets to be used are  Rotated MNIST (Larochelle et al., 2007) and Portraits over years (Ginosar et al., 2015). Indeed, it has just started to be studied in the FL community very recently. To our knowledge, there is just one work (Yoon et al., 2021) that studies FL with dynamically changing clients’ data.

Larochelle et al. An empirical evaluation of deep architectures on problems with many factors of variation. ICML 2007

Ginosar et al. A century of portraits: A visual historical record of American high school yearbooks. ICCVW 2015

Yoon et al. Federated Continual Learning with Weighted Inter-client Transfer. ICML 2021

**6. Learnable combination of $h^G$ and $h^P$.** Thanks for pointing out this interesting extension. We follow your suggestion to learn a weight $\alpha_m$ for each client $m$ to combine the two predictions by $y_P = (1-\alpha_m)h^G + \alpha_m h^P$ on CIFAR-10/100 Dir(0.3). However, we do not find a significant difference in P-FL accuracy (86.4% vs 86.1% on CIFAR-10 and 62.3% vs 62.5% on CIFAR-100). We also try learning a shared $\alpha$ for all clients but unfortunately do not see an improvement either. We surmise that when we learn $h^G+h^P$ in our current formulation, $h^P$ itself might already learn to absorb the learnable weight. That is, if we consider a similar weighted formulation, $y_P = h^G + (\alpha)/(1-\alpha) h^P$, then $h^P$ may already learn to incorporate the weight $(\alpha)/(1-\alpha)$.


 Dear reviewers and AC,

Thank you for reading our rebuttal.

We have tried to address most if not all concerns raised by the reviewers. Please let us know if you have any further questions about our paper or rebuttal. Thank you.

Best,

authors

 We thank the reviewer for the positive review and constructive comments. We provide our responses as follows.

**1. The sentence “the literature does not evaluate G-FL … local models.”** We apologize for not making it clear. What we mean is that when an FL algorithm designed for generic FL (like FedAvg) is evaluated in a personalized FL setup (i.e., evaluated on local test sets), the literature mostly applies the generic FL algorithm’s global model, rather than its local models, to the local test sets. One of our contributions is to show that the local models of generic FL algorithms are indeed strong personalized models to be evaluated on local test sets (please see the last paragraph on Page 1 and Section 3.3, and **our response to your question 4**). We will modify the footnote as “However, in the P-FL setup that every client has its own local test set, the literature does not evaluate G-FL algorithms like FedAvg by their local models but only by the global model.” in the final version.

**2&3. Averaging model weights as a regularizer.** We thank the reviewer for the comments and the references (we will cite them). First, we apologize for not making the sentence clear. What we want to claim is that taking the average over model weights indeed acts as a regularizer for “local models” to improve their individual “personalized” performance (details in Section 3.3). 

We totally agree and are aware that in machine learning, taking the average of the weights or predictions over multiple base models can improve the robustness. Indeed, in several recent works of FL (Chen & Chao, 2021; He et al., 2020; Lin et al., 2020; Zhou et al., 2020) (please see Section A.2 in the Appendix), the authors replaced weight average by bagging/model ensemble to improve the “generic” performance on the global test set. That is, they found that performing the model ensemble over clients’ models can yield more robust predictions on the global test set than the global model, which is generated by averaging the client models’ weights.

Here, we however study a different regularization effect, in personalized FL on local test sets. As reviewed in Section 3.1, personalized FL algorithms often impose a regularizer on the local/personalized models to overcome the fact that clients usually have limited data (please see Equation 3 and Equation 4 and the surrounding text). What we claim is that even without such an explicit regularizer, the model weight average before local training (Equation 2) already serves as an implicit regularizer to the local models for their individual personalized performance, as we discussed in Section 3.3 (Figure 3) and empirically verified in Section D.1 and Figure 6 in the Appendix. We hope these answer your questions and we will be happy to provide more details and discussions in the discussion phase and in the final version.


 We thank the reviewer for the positive review and constructive comments. We provide our responses as follows.

**Code and training details.** Thanks for the comments. We have included the example code in the newly uploaded supplementary materials. Due to page limits, we included the training details in the Appendix (Section C) in our initial submission, and we highlight them as follows. First, we would like to note that Fed-RoD is simple to use and requires only one additional hyperparameter compared to existing FL algorithms like FedAvg, which is the design choice of the balanced risk. For this, we simply use $\gamma=1$ for the BSM loss in Equation 6 (Ren et al., 2020) as mentioned in the paragraph “Our variants” in Section 5, following (Ren et al., 2020). The other hyperparameters for training are fixed according to FedAvg or FedDyn that we build our method upon, which include learning rate = 0.01, SGD optimizer with momentum = 0.9, local epochs = 5, weight decay = 1e-5, and batch size = 40 (please see Section C in the Appendix). 

**How many repeat trials?** We apologize for not clearly indicating them. We ran all the experiments with 5 times of different random seeds (cf. the final paragraph on Page 7). Due to the space limit, we provided the standard deviations of Table 1 in Tables 7, 8, and 9 in the Appendix (cf. the final sentence on Page 22), which are reasonably small. 

**Lack of theoretical guarantees.** Thanks for the suggestion. Having a rigorous theoretical guarantee for our work will be great, and we will consider it in our future work. We note that having it for either generic or personalized FL in non-IID conditions is by no means trivial. For instance, a concurrent work by (Cheng et al., 2021) has to make several assumptions on the non-IID distributions in order to show that fine-tuning FedAvg is a strong baseline for personalized FL (cf. Section 2). In our humble opinion, providing theoretical guarantees for our work, which deals with both generic and personalized FL, would deserve a separate paper.

That being said, our proposed work is designed upon several concepts with theoretical understanding. For instance, we incorporate the balanced risk minimization (BRM) (Section 4.1) to reduce the difference of local objectives caused by non-IID clients’ data (Section B.3 in Appendix). Our empirical results in Figure 5 justify the effectiveness of BRM, and relate it to the theoretical analysis in (Kong et al., 2021): reducing the variance of local gradients can imply better G-FL performance. We also provide an explanation of why the local models of FedAvg are strong personalized models (Section 3.3), and we empirically analyze it in section D.1 in Appendix. We also relate our work to invariant risk minimization in Section B.6 in Appendix. Overall, our work is built upon the well-justified FedAvg pipeline and we show that incorporating two losses and two prediction heads in local training leads to excellent performance for both generic and personalized FL.


 We thank the reviewer for the comments on the presentation/organization, especially in the experiment and conclusion sections. We have uploaded a revised version according to the comments. Specifically, we rewrite the experiment section by restructuring the tables and the figures according to the text, adding more discussions, and highlighting the key messages at the beginning of each paragraph. We provide a summary of our contributions in the conclusion. Due to the space limit, we have to leave some ablation studies and more experiments in the appendix but we provide a summary of them in the “More results and analyses in the Appendix” paragraph in Section 5.1. Please kindly let us know if you have further comments and we will appreciate them and be happy to incorporate them in the final version. We will also appreciate that if you would consider raising the score according to our modification.


 We thank the reviewers for their valuable comments. We are glad that the reviewers found that the problem we are solving is real and challenging in federated learning (Reviewers ZY3s, RbND); our proposed framework is “significant and novel” (Reviewer KbQt), “simple yet efficient” (Reviewer ZY3s), and “interesting” (Reviewer RbND). Both Reviewer KbQt and Reviewer ZY3s found our experiments comprehensive and extensive; Reviewer KbQt found our technical content “pretty strong”. Both Reviewer RbND and Reviewer ZY3s think the paper is easy to understand and follow. We thank Reviewer KbQt’s feedback on the presentation, and we have accordingly modified the paper in the newly updated version. In the following, we respond to the comments by each reviewer. We will incorporate all the feedback in the final version.


The paper proposes Federated Robust Decoupling (FED-ROD) to excel on both generic FL and personalized FL. With extensive results, it illustrates that strong personalized models emerge from the local training step of generic FL algorithms, and show that class-balanced objectives are effective for improving the generic FL performance when clients have different class distributions. Lastly, the paper demonstrates that FED-ROD enables zero-shot adaptation and much effective fine-tuning for new clients. The paper proposes Federated Robust Decoupling (FED-ROD) to excel on both generic FL and personalized FL. With extensive results, it illustrates that strong personalized models emerge from the local training step of generic FL algorithms, and show that class-balanced objectives are effective for improving the generic FL performance when clients have different class distributions. Lastly, the paper demonstrates that FED-ROD enables zero-shot adaptation and much effective fine-tuning for new clients.

weakness: 
1. The paper overall is a bit hard to follow. It would help to possibly re-structure tables, figures with the main body.  
2. In the experiment part, it would be stronger to expand and provide summary and insights from tables and figures in the body. For example in "Ablation Analysis",  
2. conclusion is weak and would requires some work to summarize the contribution.  Technical content of the paper is overall pretty strong and idea seems pretty significant and novel. However, overall the paper is hard to follow. It would benefit and stronger with better organization (as pointed out in the weakness in the above). 

Meta-review: This manuscript proposes and analyses can approach to address the centralized and personalized tasks in federated learning jointly. Existing work has tackled this issue by developing separate tasks. Instead, this manuscript proposes a shared architecture that aims to optimize centralized and personalized models. One observation motivating this work is that local models trained during federated learning effectively optimize local task performance. The resulting approach results well when label shifts primarily drive the client variability. Here, the centralized components are trained to optimize a balanced risk, while the local components are trained to optimize the standard empirical risk. 

Reviewers agree that the manuscript is well-written and appropriately addresses the timely issue posed. The main concerns are the clarity of the technical contributions and technical statements during the review. The authors respond to these concerns and have satisfied the reviewers. After discussion, most reviewers are generally strongly positive about the strength of the manuscript contributions.

Decision: Accept (Spotlight)